414 lines
13 KiB
Python
414 lines
13 KiB
Python
from typing import Dict, TypeVar, List, Optional
|
||
import traceback
|
||
from json_repair import repair_json
|
||
from rich.traceback import install
|
||
from src.common.logger import get_logger
|
||
from src.llm_models.utils_model import LLMRequest
|
||
from src.config.config import global_config
|
||
from src.chat.focus_chat.working_memory.memory_item import MemoryItem
|
||
import json # 添加json模块导入
|
||
|
||
|
||
install(extra_lines=3)
|
||
logger = get_logger("working_memory")
|
||
|
||
T = TypeVar("T")
|
||
|
||
|
||
class MemoryManager:
|
||
def __init__(self, chat_id: str):
|
||
"""
|
||
初始化工作记忆
|
||
|
||
Args:
|
||
chat_id: 关联的聊天ID，用于标识该工作记忆属于哪个聊天
|
||
"""
|
||
# 关联的聊天ID
|
||
self._chat_id = chat_id
|
||
|
||
# 记忆项列表
|
||
self._memories: List[MemoryItem] = []
|
||
|
||
# ID到记忆项的映射
|
||
self._id_map: Dict[str, MemoryItem] = {}
|
||
|
||
self.llm_summarizer = LLMRequest(
|
||
model=global_config.model.memory,
|
||
temperature=0.3,
|
||
request_type="working_memory",
|
||
)
|
||
|
||
@property
|
||
def chat_id(self) -> str:
|
||
"""获取关联的聊天ID"""
|
||
return self._chat_id
|
||
|
||
@chat_id.setter
|
||
def chat_id(self, value: str):
|
||
"""设置关联的聊天ID"""
|
||
self._chat_id = value
|
||
|
||
def push_item(self, memory_item: MemoryItem) -> str:
|
||
"""
|
||
推送一个已创建的记忆项到工作记忆中
|
||
|
||
Args:
|
||
memory_item: 要存储的记忆项
|
||
|
||
Returns:
|
||
记忆项的ID
|
||
"""
|
||
# 添加到内存和ID映射
|
||
self._memories.append(memory_item)
|
||
self._id_map[memory_item.id] = memory_item
|
||
|
||
return memory_item.id
|
||
|
||
def get_by_id(self, memory_id: str) -> Optional[MemoryItem]:
|
||
"""
|
||
通过ID获取记忆项
|
||
|
||
Args:
|
||
memory_id: 记忆项ID
|
||
|
||
Returns:
|
||
找到的记忆项，如果不存在则返回None
|
||
"""
|
||
memory_item = self._id_map.get(memory_id)
|
||
if memory_item:
|
||
# 检查记忆强度，如果小于1则删除
|
||
if not memory_item.is_memory_valid():
|
||
print(f"记忆 {memory_id} 强度过低 ({memory_item.memory_strength})，已自动移除")
|
||
self.delete(memory_id)
|
||
return None
|
||
|
||
return memory_item
|
||
|
||
def get_all_items(self) -> List[MemoryItem]:
|
||
"""获取所有记忆项"""
|
||
return list(self._id_map.values())
|
||
|
||
def find_items(
|
||
self,
|
||
source: Optional[str] = None,
|
||
start_time: Optional[float] = None,
|
||
end_time: Optional[float] = None,
|
||
memory_id: Optional[str] = None,
|
||
limit: Optional[int] = None,
|
||
newest_first: bool = False,
|
||
min_strength: float = 0.0,
|
||
) -> List[MemoryItem]:
|
||
"""
|
||
按条件查找记忆项
|
||
|
||
Args:
|
||
source: 数据来源
|
||
start_time: 开始时间戳
|
||
end_time: 结束时间戳
|
||
memory_id: 特定记忆项ID
|
||
limit: 返回结果的最大数量
|
||
newest_first: 是否按最新优先排序
|
||
min_strength: 最小记忆强度
|
||
|
||
Returns:
|
||
符合条件的记忆项列表
|
||
"""
|
||
# 如果提供了特定ID，直接查找
|
||
if memory_id:
|
||
item = self.get_by_id(memory_id)
|
||
return [item] if item else []
|
||
|
||
results = []
|
||
|
||
# 获取所有项目
|
||
items = self._memories
|
||
|
||
# 如果需要最新优先，则反转遍历顺序
|
||
if newest_first:
|
||
items_to_check = list(reversed(items))
|
||
else:
|
||
items_to_check = items
|
||
|
||
# 遍历项目
|
||
for item in items_to_check:
|
||
# 检查来源是否匹配
|
||
if source is not None and not item.matches_source(source):
|
||
continue
|
||
|
||
# 检查时间范围
|
||
if start_time is not None and item.timestamp < start_time:
|
||
continue
|
||
if end_time is not None and item.timestamp > end_time:
|
||
continue
|
||
|
||
# 检查记忆强度
|
||
if min_strength > 0 and item.memory_strength < min_strength:
|
||
continue
|
||
|
||
# 所有条件都满足，添加到结果中
|
||
results.append(item)
|
||
|
||
# 如果达到限制数量，提前返回
|
||
if limit is not None and len(results) >= limit:
|
||
return results
|
||
|
||
return results
|
||
|
||
async def summarize_memory_item(self, content: str) -> Dict[str, str]:
|
||
"""
|
||
使用LLM总结记忆项
|
||
|
||
Args:
|
||
content: 需要总结的内容
|
||
|
||
Returns:
|
||
包含brief和summary的字典
|
||
"""
|
||
prompt = f"""请对以下内容进行总结，总结成记忆，输出两部分：
|
||
1. 记忆内容主题（精简，20字以内）：让用户可以一眼看出记忆内容是什么
|
||
2. 记忆内容概括：对内容进行概括，保留重要信息，200字以内
|
||
|
||
内容：
|
||
{content}
|
||
|
||
请按以下JSON格式输出：
|
||
{{
|
||
"brief": "记忆内容主题",
|
||
"summary": "记忆内容概括"
|
||
}}
|
||
请确保输出是有效的JSON格式，不要添加任何额外的说明或解释。
|
||
"""
|
||
default_summary = {
|
||
"brief": "主题未知的记忆",
|
||
"summary": "无法概括的记忆内容",
|
||
}
|
||
|
||
try:
|
||
# 调用LLM生成总结
|
||
response, _ = await self.llm_summarizer.generate_response_async(prompt)
|
||
|
||
# 使用repair_json解析响应
|
||
try:
|
||
# 使用repair_json修复JSON格式
|
||
fixed_json_string = repair_json(response)
|
||
|
||
# 如果repair_json返回的是字符串，需要解析为Python对象
|
||
if isinstance(fixed_json_string, str):
|
||
try:
|
||
json_result = json.loads(fixed_json_string)
|
||
except json.JSONDecodeError as decode_error:
|
||
logger.error(f"JSON解析错误: {str(decode_error)}")
|
||
return default_summary
|
||
else:
|
||
# 如果repair_json直接返回了字典对象，直接使用
|
||
json_result = fixed_json_string
|
||
|
||
# 进行额外的类型检查
|
||
if not isinstance(json_result, dict):
|
||
logger.error(f"修复后的JSON不是字典类型: {type(json_result)}")
|
||
return default_summary
|
||
|
||
# 确保所有必要字段都存在且类型正确
|
||
if "brief" not in json_result or not isinstance(json_result["brief"], str):
|
||
json_result["brief"] = "主题未知的记忆"
|
||
|
||
if "summary" not in json_result or not isinstance(json_result["summary"], str):
|
||
json_result["summary"] = "无法概括的记忆内容"
|
||
|
||
return json_result
|
||
|
||
except Exception as json_error:
|
||
logger.error(f"JSON处理失败: {str(json_error)}，将使用默认摘要")
|
||
return default_summary
|
||
|
||
except Exception as e:
|
||
logger.error(f"生成总结时出错: {str(e)}")
|
||
return default_summary
|
||
|
||
def decay_memory(self, memory_id: str, decay_factor: float = 0.8) -> bool:
|
||
"""
|
||
使单个记忆衰减
|
||
|
||
Args:
|
||
memory_id: 记忆ID
|
||
decay_factor: 衰减因子(0-1之间)
|
||
|
||
Returns:
|
||
是否成功衰减
|
||
"""
|
||
memory_item = self.get_by_id(memory_id)
|
||
if not memory_item:
|
||
return False
|
||
|
||
# 计算衰减量（当前强度 * (1-衰减因子)）
|
||
old_strength = memory_item.memory_strength
|
||
decay_amount = old_strength * (1 - decay_factor)
|
||
|
||
# 更新强度
|
||
memory_item.memory_strength = decay_amount
|
||
|
||
return True
|
||
|
||
def delete(self, memory_id: str) -> bool:
|
||
"""
|
||
删除指定ID的记忆项
|
||
|
||
Args:
|
||
memory_id: 要删除的记忆项ID
|
||
|
||
Returns:
|
||
是否成功删除
|
||
"""
|
||
if memory_id not in self._id_map:
|
||
return False
|
||
|
||
# 获取要删除的项
|
||
self._id_map[memory_id]
|
||
|
||
# 从内存中删除
|
||
self._memories = [i for i in self._memories if i.id != memory_id]
|
||
|
||
# 从ID映射中删除
|
||
del self._id_map[memory_id]
|
||
|
||
return True
|
||
|
||
def clear(self) -> None:
|
||
"""清除所有记忆"""
|
||
self._memories.clear()
|
||
self._id_map.clear()
|
||
|
||
async def merge_memories(
|
||
self, memory_id1: str, memory_id2: str, reason: str, delete_originals: bool = True
|
||
) -> MemoryItem:
|
||
"""
|
||
合并两个记忆项
|
||
|
||
Args:
|
||
memory_id1: 第一个记忆项ID
|
||
memory_id2: 第二个记忆项ID
|
||
reason: 合并原因
|
||
delete_originals: 是否删除原始记忆，默认为True
|
||
|
||
Returns:
|
||
合并后的记忆项
|
||
"""
|
||
# 获取两个记忆项
|
||
memory_item1 = self.get_by_id(memory_id1)
|
||
memory_item2 = self.get_by_id(memory_id2)
|
||
|
||
if not memory_item1 or not memory_item2:
|
||
raise ValueError("无法找到指定的记忆项")
|
||
|
||
# 构建合并提示
|
||
prompt = f"""
|
||
请根据以下原因，将两段记忆内容有机合并成一段新的记忆内容。
|
||
合并时保留两段记忆的重要信息，避免重复，确保生成的内容连贯、自然。
|
||
|
||
合并原因：{reason}
|
||
|
||
记忆1主题：{memory_item1.brief}
|
||
记忆1内容：{memory_item1.summary}
|
||
|
||
记忆2主题：{memory_item2.brief}
|
||
记忆2内容：{memory_item2.summary}
|
||
|
||
请按以下JSON格式输出合并结果：
|
||
{{
|
||
"brief": "合并后的主题（20字以内）",
|
||
"summary": "合并后的内容概括（200字以内）"
|
||
}}
|
||
请确保输出是有效的JSON格式，不要添加任何额外的说明或解释。
|
||
"""
|
||
|
||
# 默认合并结果
|
||
default_merged = {
|
||
"brief": f"合并：{memory_item1.brief} + {memory_item2.brief}",
|
||
"summary": f"合并的记忆：{memory_item1.summary}\n{memory_item2.summary}",
|
||
}
|
||
|
||
try:
|
||
# 调用LLM合并记忆
|
||
response, _ = await self.llm_summarizer.generate_response_async(prompt)
|
||
|
||
# 处理LLM返回的合并结果
|
||
try:
|
||
# 修复JSON格式
|
||
fixed_json_string = repair_json(response)
|
||
|
||
# 将修复后的字符串解析为Python对象
|
||
if isinstance(fixed_json_string, str):
|
||
try:
|
||
merged_data = json.loads(fixed_json_string)
|
||
except json.JSONDecodeError as decode_error:
|
||
logger.error(f"JSON解析错误: {str(decode_error)}")
|
||
merged_data = default_merged
|
||
else:
|
||
# 如果repair_json直接返回了字典对象，直接使用
|
||
merged_data = fixed_json_string
|
||
|
||
# 确保是字典类型
|
||
if not isinstance(merged_data, dict):
|
||
logger.error(f"修复后的JSON不是字典类型: {type(merged_data)}")
|
||
merged_data = default_merged
|
||
|
||
if "brief" not in merged_data or not isinstance(merged_data["brief"], str):
|
||
merged_data["brief"] = default_merged["brief"]
|
||
|
||
if "summary" not in merged_data or not isinstance(merged_data["summary"], str):
|
||
merged_data["summary"] = default_merged["summary"]
|
||
|
||
except Exception as e:
|
||
logger.error(f"合并记忆时处理JSON出错: {str(e)}")
|
||
traceback.print_exc()
|
||
merged_data = default_merged
|
||
except Exception as e:
|
||
logger.error(f"合并记忆调用LLM出错: {str(e)}")
|
||
traceback.print_exc()
|
||
merged_data = default_merged
|
||
|
||
# 创建新的记忆项
|
||
# 取两个记忆项中更强的来源
|
||
merged_source = (
|
||
memory_item1.from_source
|
||
if memory_item1.memory_strength >= memory_item2.memory_strength
|
||
else memory_item2.from_source
|
||
)
|
||
|
||
# 创建新的记忆项
|
||
merged_memory = MemoryItem(
|
||
summary=merged_data["summary"], from_source=merged_source, brief=merged_data["brief"]
|
||
)
|
||
|
||
# 记忆强度取两者最大值
|
||
merged_memory.memory_strength = max(memory_item1.memory_strength, memory_item2.memory_strength)
|
||
|
||
# 添加到存储中
|
||
self.push_item(merged_memory)
|
||
|
||
# 如果需要，删除原始记忆
|
||
if delete_originals:
|
||
self.delete(memory_id1)
|
||
self.delete(memory_id2)
|
||
|
||
return merged_memory
|
||
|
||
def delete_earliest_memory(self) -> bool:
|
||
"""
|
||
删除最早的记忆项
|
||
|
||
Returns:
|
||
是否成功删除
|
||
"""
|
||
# 获取所有记忆项
|
||
all_memories = self.get_all_items()
|
||
|
||
if not all_memories:
|
||
return False
|
||
|
||
# 按时间戳排序，找到最早的记忆项
|
||
earliest_memory = min(all_memories, key=lambda item: item.timestamp)
|
||
|
||
# 删除最早的记忆项
|
||
return self.delete(earliest_memory.id)
|