460 lines
18 KiB
Python
460 lines
18 KiB
Python
# -*- coding: utf-8 -*-
|
||
import sys
|
||
import jieba
|
||
import networkx as nx
|
||
import matplotlib.pyplot as plt
|
||
import math
|
||
from collections import Counter
|
||
import datetime
|
||
import random
|
||
import time
|
||
import os
|
||
# from chat.config import global_config
|
||
sys.path.append("C:/GitHub/MaiMBot") # 添加项目根目录到 Python 路径
|
||
from src.common.database import Database # 使用正确的导入语法
|
||
from src.plugins.memory_system.llm_module import LLMModel
|
||
|
||
def calculate_information_content(text):
|
||
"""计算文本的信息量（熵）"""
|
||
# 统计字符频率
|
||
char_count = Counter(text)
|
||
total_chars = len(text)
|
||
|
||
# 计算熵
|
||
entropy = 0
|
||
for count in char_count.values():
|
||
probability = count / total_chars
|
||
entropy -= probability * math.log2(probability)
|
||
|
||
return entropy
|
||
|
||
def get_cloest_chat_from_db(db, length: int, timestamp: str):
|
||
"""从数据库中获取最接近指定时间戳的聊天记录"""
|
||
chat_text = ''
|
||
closest_record = db.db.messages.find_one({"time": {"$lte": timestamp}}, sort=[('time', -1)])
|
||
|
||
if closest_record:
|
||
closest_time = closest_record['time']
|
||
group_id = closest_record['group_id'] # 获取groupid
|
||
# 获取该时间戳之后的length条消息，且groupid相同
|
||
chat_record = list(db.db.messages.find({"time": {"$gt": closest_time}, "group_id": group_id}).sort('time', 1).limit(length))
|
||
for record in chat_record:
|
||
time_str = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(int(record['time'])))
|
||
chat_text += f'[{time_str}] {record["user_nickname"] or "用户" + str(record["user_id"])}: {record["processed_plain_text"]}\n'
|
||
return chat_text
|
||
|
||
return ''
|
||
|
||
class Memory_graph:
|
||
def __init__(self):
|
||
self.G = nx.Graph() # 使用 networkx 的图结构
|
||
self.db = Database.get_instance()
|
||
|
||
def connect_dot(self, concept1, concept2):
|
||
self.G.add_edge(concept1, concept2)
|
||
|
||
def add_dot(self, concept, memory):
|
||
if concept in self.G:
|
||
# 如果节点已存在，将新记忆添加到现有列表中
|
||
if 'memory_items' in self.G.nodes[concept]:
|
||
if not isinstance(self.G.nodes[concept]['memory_items'], list):
|
||
# 如果当前不是列表，将其转换为列表
|
||
self.G.nodes[concept]['memory_items'] = [self.G.nodes[concept]['memory_items']]
|
||
self.G.nodes[concept]['memory_items'].append(memory)
|
||
else:
|
||
self.G.nodes[concept]['memory_items'] = [memory]
|
||
else:
|
||
# 如果是新节点，创建新的记忆列表
|
||
self.G.add_node(concept, memory_items=[memory])
|
||
|
||
def get_dot(self, concept):
|
||
# 检查节点是否存在于图中
|
||
if concept in self.G:
|
||
# 从图中获取节点数据
|
||
node_data = self.G.nodes[concept]
|
||
# print(node_data)
|
||
# 创建新的Memory_dot对象
|
||
return concept,node_data
|
||
return None
|
||
|
||
def get_related_item(self, topic, depth=1):
|
||
if topic not in self.G:
|
||
return [], []
|
||
|
||
first_layer_items = []
|
||
second_layer_items = []
|
||
|
||
# 获取相邻节点
|
||
neighbors = list(self.G.neighbors(topic))
|
||
# print(f"第一层: {topic}")
|
||
|
||
# 获取当前节点的记忆项
|
||
node_data = self.get_dot(topic)
|
||
if node_data:
|
||
concept, data = node_data
|
||
if 'memory_items' in data:
|
||
memory_items = data['memory_items']
|
||
if isinstance(memory_items, list):
|
||
first_layer_items.extend(memory_items)
|
||
else:
|
||
first_layer_items.append(memory_items)
|
||
|
||
# 只在depth=2时获取第二层记忆
|
||
if depth >= 2:
|
||
# 获取相邻节点的记忆项
|
||
for neighbor in neighbors:
|
||
# print(f"第二层: {neighbor}")
|
||
node_data = self.get_dot(neighbor)
|
||
if node_data:
|
||
concept, data = node_data
|
||
if 'memory_items' in data:
|
||
memory_items = data['memory_items']
|
||
if isinstance(memory_items, list):
|
||
second_layer_items.extend(memory_items)
|
||
else:
|
||
second_layer_items.append(memory_items)
|
||
|
||
return first_layer_items, second_layer_items
|
||
|
||
def store_memory(self):
|
||
for node in self.G.nodes():
|
||
dot_data = {
|
||
"concept": node
|
||
}
|
||
self.db.db.store_memory_dots.insert_one(dot_data)
|
||
|
||
@property
|
||
def dots(self):
|
||
# 返回所有节点对应的 Memory_dot 对象
|
||
return [self.get_dot(node) for node in self.G.nodes()]
|
||
|
||
|
||
def get_random_chat_from_db(self, length: int, timestamp: str):
|
||
# 从数据库中根据时间戳获取离其最近的聊天记录
|
||
chat_text = ''
|
||
closest_record = self.db.db.messages.find_one({"time": {"$lte": timestamp}}, sort=[('time', -1)]) # 调试输出
|
||
|
||
# print(f"距离time最近的消息时间: {time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(int(closest_record['time'])))}")
|
||
|
||
if closest_record:
|
||
closest_time = closest_record['time']
|
||
group_id = closest_record['group_id'] # 获取groupid
|
||
# 获取该时间戳之后的length条消息，且groupid相同
|
||
chat_record = list(self.db.db.messages.find({"time": {"$gt": closest_time}, "group_id": group_id}).sort('time', 1).limit(length))
|
||
for record in chat_record:
|
||
if record:
|
||
time_str = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(int(record['time'])))
|
||
chat_text += f'[{time_str}] {record["user_nickname"] or "用户" + str(record["user_id"])}: {record["processed_plain_text"]}\n' # 添加发送者和时间信息
|
||
return chat_text
|
||
|
||
return [] # 如果没有找到记录，返回空列表
|
||
|
||
def save_graph_to_db(self):
|
||
# 保存节点
|
||
for node in self.G.nodes(data=True):
|
||
concept = node[0]
|
||
memory_items = node[1].get('memory_items', [])
|
||
|
||
# 查找是否存在同名节点
|
||
existing_node = self.db.db.graph_data.nodes.find_one({'concept': concept})
|
||
if existing_node:
|
||
# 如果存在,合并memory_items并去重
|
||
existing_items = existing_node.get('memory_items', [])
|
||
if not isinstance(existing_items, list):
|
||
existing_items = [existing_items] if existing_items else []
|
||
|
||
# 合并并去重
|
||
all_items = list(set(existing_items + memory_items))
|
||
|
||
# 更新节点
|
||
self.db.db.graph_data.nodes.update_one(
|
||
{'concept': concept},
|
||
{'$set': {'memory_items': all_items}}
|
||
)
|
||
else:
|
||
# 如果不存在,创建新节点
|
||
node_data = {
|
||
'concept': concept,
|
||
'memory_items': memory_items
|
||
}
|
||
self.db.db.graph_data.nodes.insert_one(node_data)
|
||
|
||
# 保存边
|
||
for edge in self.G.edges():
|
||
source, target = edge
|
||
|
||
# 查找是否存在同样的边
|
||
existing_edge = self.db.db.graph_data.edges.find_one({
|
||
'source': source,
|
||
'target': target
|
||
})
|
||
|
||
if existing_edge:
|
||
# 如果存在,增加num属性
|
||
num = existing_edge.get('num', 1) + 1
|
||
self.db.db.graph_data.edges.update_one(
|
||
{'source': source, 'target': target},
|
||
{'$set': {'num': num}}
|
||
)
|
||
else:
|
||
# 如果不存在,创建新边
|
||
edge_data = {
|
||
'source': source,
|
||
'target': target,
|
||
'num': 1
|
||
}
|
||
self.db.db.graph_data.edges.insert_one(edge_data)
|
||
|
||
def load_graph_from_db(self):
|
||
# 清空当前图
|
||
self.G.clear()
|
||
# 加载节点
|
||
nodes = self.db.db.graph_data.nodes.find()
|
||
for node in nodes:
|
||
memory_items = node.get('memory_items', [])
|
||
if not isinstance(memory_items, list):
|
||
memory_items = [memory_items] if memory_items else []
|
||
self.G.add_node(node['concept'], memory_items=memory_items)
|
||
# 加载边
|
||
edges = self.db.db.graph_data.edges.find()
|
||
for edge in edges:
|
||
self.G.add_edge(edge['source'], edge['target'], num=edge.get('num', 1))
|
||
|
||
# 海马体
|
||
class Hippocampus:
|
||
def __init__(self,memory_graph:Memory_graph):
|
||
self.memory_graph = memory_graph
|
||
self.llm_model = LLMModel()
|
||
self.llm_model_small = LLMModel(model_name="deepseek-ai/DeepSeek-V2.5")
|
||
|
||
def get_memory_sample(self,chat_size=20,time_frequency:dict={'near':2,'mid':4,'far':3}):
|
||
current_timestamp = datetime.datetime.now().timestamp()
|
||
chat_text = []
|
||
#短期：1h 中期：4h 长期：24h
|
||
for _ in range(time_frequency.get('near')): # 循环10次
|
||
random_time = current_timestamp - random.randint(1, 3600) # 随机时间
|
||
chat_ = get_cloest_chat_from_db(db=self.memory_graph.db, length=chat_size, timestamp=random_time)
|
||
chat_text.append(chat_)
|
||
for _ in range(time_frequency.get('mid')): # 循环10次
|
||
random_time = current_timestamp - random.randint(3600, 3600*4) # 随机时间
|
||
chat_ = get_cloest_chat_from_db(db=self.memory_graph.db, length=chat_size, timestamp=random_time)
|
||
chat_text.append(chat_)
|
||
for _ in range(time_frequency.get('far')): # 循环10次
|
||
random_time = current_timestamp - random.randint(3600*4, 3600*24) # 随机时间
|
||
chat_ = get_cloest_chat_from_db(db=self.memory_graph.db, length=chat_size, timestamp=random_time)
|
||
chat_text.append(chat_)
|
||
return chat_text
|
||
|
||
def build_memory(self,chat_size=12):
|
||
#最近消息获取频率
|
||
time_frequency = {'near':1,'mid':2,'far':2}
|
||
memory_sample = self.get_memory_sample(chat_size,time_frequency)
|
||
|
||
#加载进度可视化
|
||
for i, input_text in enumerate(memory_sample, 1):
|
||
progress = (i / len(memory_sample)) * 100
|
||
bar_length = 30
|
||
filled_length = int(bar_length * i // len(memory_sample))
|
||
bar = '█' * filled_length + '-' * (bar_length - filled_length)
|
||
print(f"\n进度: [{bar}] {progress:.1f}% ({i}/{len(memory_sample)})")
|
||
# print(f"第{i}条消息: {input_text}")
|
||
if input_text:
|
||
# 生成压缩后记忆
|
||
first_memory = set()
|
||
first_memory = self.memory_compress(input_text, 2.5)
|
||
#将记忆加入到图谱中
|
||
for topic, memory in first_memory:
|
||
topics = segment_text(topic)
|
||
print(f"\033[1;34m话题\033[0m: {topic},节点: {topics}, 记忆: {memory}")
|
||
for split_topic in topics:
|
||
self.memory_graph.add_dot(split_topic,memory)
|
||
for split_topic in topics:
|
||
for other_split_topic in topics:
|
||
if split_topic != other_split_topic:
|
||
self.memory_graph.connect_dot(split_topic, other_split_topic)
|
||
else:
|
||
print(f"空消息 跳过")
|
||
|
||
self.memory_graph.save_graph_to_db()
|
||
|
||
def memory_compress(self, input_text, rate=1):
|
||
information_content = calculate_information_content(input_text)
|
||
print(f"文本的信息量（熵）: {information_content:.4f} bits")
|
||
topic_num = max(1, min(5, int(information_content * rate / 4)))
|
||
topic_prompt = find_topic(input_text, topic_num)
|
||
topic_response = self.llm_model.generate_response(topic_prompt)
|
||
# 检查 topic_response 是否为元组
|
||
if isinstance(topic_response, tuple):
|
||
topics = topic_response[0].split(",") # 假设第一个元素是我们需要的字符串
|
||
else:
|
||
topics = topic_response.split(",")
|
||
compressed_memory = set()
|
||
for topic in topics:
|
||
topic_what_prompt = topic_what(input_text,topic)
|
||
topic_what_response = self.llm_model_small.generate_response(topic_what_prompt)
|
||
compressed_memory.add((topic.strip(), topic_what_response[0])) # 将话题和记忆作为元组存储
|
||
return compressed_memory
|
||
|
||
def segment_text(text):
|
||
seg_text = list(jieba.cut(text))
|
||
return seg_text
|
||
|
||
def find_topic(text, topic_num):
|
||
prompt = f'这是一段文字：{text}。请你从这段话中总结出{topic_num}个话题，帮我列出来，用逗号隔开，尽可能精简。只需要列举{topic_num}个话题就好，不要告诉我其他内容。'
|
||
return prompt
|
||
|
||
def topic_what(text, topic):
|
||
prompt = f'这是一段文字：{text}。我想知道这记忆里有什么关于{topic}的话题，帮我总结成一句自然的话，可以包含时间和人物。只输出这句话就好'
|
||
return prompt
|
||
|
||
def visualize_graph(memory_graph: Memory_graph, color_by_memory: bool = False):
|
||
# 设置中文字体
|
||
plt.rcParams['font.sans-serif'] = ['SimHei'] # 用来正常显示中文标签
|
||
plt.rcParams['axes.unicode_minus'] = False # 用来正常显示负号
|
||
|
||
G = memory_graph.G
|
||
|
||
# 创建一个新图用于可视化
|
||
H = G.copy()
|
||
|
||
# 移除只有一条记忆的节点和连接数少于3的节点
|
||
nodes_to_remove = []
|
||
for node in H.nodes():
|
||
memory_items = H.nodes[node].get('memory_items', [])
|
||
memory_count = len(memory_items) if isinstance(memory_items, list) else (1 if memory_items else 0)
|
||
degree = H.degree(node)
|
||
if memory_count <= 1 or degree <= 2:
|
||
nodes_to_remove.append(node)
|
||
|
||
H.remove_nodes_from(nodes_to_remove)
|
||
|
||
# 如果过滤后没有节点，则返回
|
||
if len(H.nodes()) == 0:
|
||
print("过滤后没有符合条件的节点可显示")
|
||
return
|
||
|
||
# 保存图到本地
|
||
nx.write_gml(H, "memory_graph.gml") # 保存为 GML 格式
|
||
|
||
# 根据连接条数或记忆数量设置节点颜色
|
||
node_colors = []
|
||
nodes = list(H.nodes()) # 获取图中实际的节点列表
|
||
|
||
if color_by_memory:
|
||
# 计算每个节点的记忆数量
|
||
memory_counts = []
|
||
for node in nodes:
|
||
memory_items = H.nodes[node].get('memory_items', [])
|
||
if isinstance(memory_items, list):
|
||
count = len(memory_items)
|
||
else:
|
||
count = 1 if memory_items else 0
|
||
memory_counts.append(count)
|
||
max_memories = max(memory_counts) if memory_counts else 1
|
||
|
||
for count in memory_counts:
|
||
# 使用不同的颜色方案：红色表示记忆多，蓝色表示记忆少
|
||
if max_memories > 0:
|
||
intensity = min(1.0, count / max_memories)
|
||
color = (intensity, 0, 1.0 - intensity) # 从蓝色渐变到红色
|
||
else:
|
||
color = (0, 0, 1) # 如果没有记忆，则为蓝色
|
||
node_colors.append(color)
|
||
else:
|
||
# 使用原来的连接数量着色方案
|
||
max_degree = max(H.degree(), key=lambda x: x[1])[1] if H.degree() else 1
|
||
for node in nodes:
|
||
degree = H.degree(node)
|
||
if max_degree > 0:
|
||
red = min(1.0, degree / max_degree)
|
||
blue = 1.0 - red
|
||
color = (red, 0, blue)
|
||
else:
|
||
color = (0, 0, 1)
|
||
node_colors.append(color)
|
||
|
||
# 绘制图形
|
||
plt.figure(figsize=(12, 8))
|
||
pos = nx.spring_layout(H, k=1, iterations=50)
|
||
nx.draw(H, pos,
|
||
with_labels=True,
|
||
node_color=node_colors,
|
||
node_size=2000,
|
||
font_size=10,
|
||
font_family='SimHei',
|
||
font_weight='bold')
|
||
|
||
title = '记忆图谱可视化 - ' + ('按记忆数量着色' if color_by_memory else '按连接数量着色')
|
||
plt.title(title, fontsize=16, fontfamily='SimHei')
|
||
plt.show()
|
||
|
||
def main():
|
||
# 初始化数据库
|
||
Database.initialize(
|
||
host= os.getenv("MONGODB_HOST"),
|
||
port= int(os.getenv("MONGODB_PORT")),
|
||
db_name= os.getenv("DATABASE_NAME"),
|
||
username= os.getenv("MONGODB_USERNAME"),
|
||
password= os.getenv("MONGODB_PASSWORD"),
|
||
auth_source=os.getenv("MONGODB_AUTH_SOURCE")
|
||
)
|
||
|
||
start_time = time.time()
|
||
|
||
# 创建记忆图
|
||
memory_graph = Memory_graph()
|
||
# 加载数据库中存储的记忆图
|
||
memory_graph.load_graph_from_db()
|
||
# 创建海马体
|
||
hippocampus = Hippocampus(memory_graph)
|
||
|
||
end_time = time.time()
|
||
print(f"\033[32m[加载海马体耗时: {end_time - start_time:.2f} 秒]\033[0m")
|
||
|
||
# 构建记忆
|
||
hippocampus.build_memory(chat_size=25)
|
||
|
||
# 展示两种不同的可视化方式
|
||
print("\n按连接数量着色的图谱：")
|
||
visualize_graph(memory_graph, color_by_memory=False)
|
||
|
||
print("\n按记忆数量着色的图谱：")
|
||
visualize_graph(memory_graph, color_by_memory=True)
|
||
|
||
# 交互式查询
|
||
while True:
|
||
query = input("请输入新的查询概念（输入'退出'以结束）：")
|
||
if query.lower() == '退出':
|
||
break
|
||
items_list = memory_graph.get_related_item(query)
|
||
if items_list:
|
||
for memory_item in items_list:
|
||
print(memory_item)
|
||
else:
|
||
print("未找到相关记忆。")
|
||
|
||
while True:
|
||
query = input("请输入问题：")
|
||
|
||
if query.lower() == '退出':
|
||
break
|
||
|
||
topic_prompt = find_topic(query, 3)
|
||
topic_response = hippocampus.llm_model.generate_response(topic_prompt)
|
||
# 检查 topic_response 是否为元组
|
||
if isinstance(topic_response, tuple):
|
||
topics = topic_response[0].split(",") # 假设第一个元素是我们需要的字符串
|
||
else:
|
||
topics = topic_response.split(",")
|
||
print(topics)
|
||
|
||
for keyword in topics:
|
||
items_list = memory_graph.get_related_item(keyword)
|
||
if items_list:
|
||
print(items_list)
|
||
|
||
if __name__ == "__main__":
|
||
main()
|
||
|
||
|