707 lines
28 KiB
Python
707 lines
28 KiB
Python
"""
|
||
统一记忆管理器 (Unified Memory Manager)
|
||
|
||
整合三层记忆系统：
|
||
- 感知记忆层
|
||
- 短期记忆层
|
||
- 长期记忆层
|
||
|
||
提供统一的接口供外部调用
|
||
"""
|
||
|
||
import asyncio
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
from src.common.logger import get_logger
|
||
from src.memory_graph.long_term_manager import LongTermMemoryManager
|
||
from src.memory_graph.manager import MemoryManager
|
||
from src.memory_graph.models import JudgeDecision, MemoryBlock, ShortTermMemory
|
||
from src.memory_graph.perceptual_manager import PerceptualMemoryManager
|
||
from src.memory_graph.short_term_manager import ShortTermMemoryManager
|
||
|
||
logger = get_logger(__name__)
|
||
|
||
|
||
class UnifiedMemoryManager:
|
||
"""
|
||
统一记忆管理器
|
||
|
||
整合三层记忆系统，提供统一接口
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
data_dir: Path | None = None,
|
||
memory_manager: MemoryManager | None = None,
|
||
# 感知记忆配置
|
||
perceptual_max_blocks: int = 50,
|
||
perceptual_block_size: int = 5,
|
||
perceptual_activation_threshold: int = 3,
|
||
perceptual_recall_top_k: int = 5,
|
||
perceptual_recall_threshold: float = 0.55,
|
||
# 短期记忆配置
|
||
short_term_max_memories: int = 30,
|
||
short_term_transfer_threshold: float = 0.6,
|
||
short_term_overflow_strategy: str = "transfer_all",
|
||
short_term_enable_force_cleanup: bool = False,
|
||
short_term_cleanup_keep_ratio: float = 0.9,
|
||
# 长期记忆配置
|
||
long_term_batch_size: int = 10,
|
||
long_term_search_top_k: int = 5,
|
||
long_term_decay_factor: float = 0.95,
|
||
long_term_auto_transfer_interval: int = 600,
|
||
# 智能检索配置
|
||
judge_confidence_threshold: float = 0.7,
|
||
):
|
||
"""
|
||
初始化统一记忆管理器
|
||
|
||
Args:
|
||
data_dir: 数据存储目录
|
||
perceptual_max_blocks: 感知记忆堆最大容量
|
||
perceptual_block_size: 每个记忆块的消息数量
|
||
perceptual_activation_threshold: 激活阈值（召回次数）
|
||
perceptual_recall_top_k: 召回时返回的最大块数
|
||
perceptual_recall_threshold: 召回的相似度阈值
|
||
short_term_max_memories: 短期记忆最大数量
|
||
short_term_transfer_threshold: 转移到长期记忆的重要性阈值
|
||
long_term_batch_size: 批量处理的短期记忆数量
|
||
long_term_search_top_k: 检索相似记忆的数量
|
||
long_term_decay_factor: 长期记忆的衰减因子
|
||
long_term_auto_transfer_interval: 自动转移间隔（秒）
|
||
judge_confidence_threshold: 裁判模型的置信度阈值
|
||
"""
|
||
self.data_dir = data_dir or Path("data/memory_graph")
|
||
self.data_dir.mkdir(parents=True, exist_ok=True)
|
||
|
||
# 配置参数
|
||
self.judge_confidence_threshold = judge_confidence_threshold
|
||
|
||
# 三层管理器
|
||
self.perceptual_manager: PerceptualMemoryManager
|
||
self.short_term_manager: ShortTermMemoryManager
|
||
self.long_term_manager: LongTermMemoryManager
|
||
|
||
# 底层 MemoryManager（长期记忆）
|
||
self.memory_manager: MemoryManager = memory_manager
|
||
|
||
# 配置参数存储（用于初始化）
|
||
self._config = {
|
||
"perceptual": {
|
||
"max_blocks": perceptual_max_blocks,
|
||
"block_size": perceptual_block_size,
|
||
"activation_threshold": perceptual_activation_threshold,
|
||
"recall_top_k": perceptual_recall_top_k,
|
||
"recall_similarity_threshold": perceptual_recall_threshold,
|
||
},
|
||
"short_term": {
|
||
"max_memories": short_term_max_memories,
|
||
"transfer_importance_threshold": short_term_transfer_threshold,
|
||
"overflow_strategy": short_term_overflow_strategy,
|
||
"enable_force_cleanup": short_term_enable_force_cleanup,
|
||
"cleanup_keep_ratio": short_term_cleanup_keep_ratio,
|
||
},
|
||
"long_term": {
|
||
"batch_size": long_term_batch_size,
|
||
"search_top_k": long_term_search_top_k,
|
||
"long_term_decay_factor": long_term_decay_factor,
|
||
},
|
||
}
|
||
|
||
# 状态
|
||
self._initialized = False
|
||
self._auto_transfer_task: asyncio.Task | None = None
|
||
self._auto_transfer_interval = max(10.0, float(long_term_auto_transfer_interval))
|
||
self._transfer_wakeup_event: asyncio.Event | None = None
|
||
|
||
logger.info("统一记忆管理器已创建")
|
||
|
||
async def initialize(self) -> None:
|
||
"""初始化统一记忆管理器"""
|
||
if self._initialized:
|
||
logger.warning("统一记忆管理器已经初始化")
|
||
return
|
||
|
||
try:
|
||
logger.debug("开始初始化统一记忆管理器...")
|
||
|
||
# 初始化底层 MemoryManager（长期记忆）
|
||
if self.memory_manager is None:
|
||
# 如果未提供外部 MemoryManager，则创建一个新的
|
||
# 假设 data_dir 是 three_tier 子目录，则 MemoryManager 使用父目录
|
||
# 如果 data_dir 是根目录，则 MemoryManager 使用该目录
|
||
self.memory_manager = MemoryManager(data_dir=self.data_dir)
|
||
await self.memory_manager.initialize()
|
||
else:
|
||
logger.debug("使用外部提供的 MemoryManager")
|
||
# 确保外部 MemoryManager 已初始化
|
||
if not getattr(self.memory_manager, "_initialized", False):
|
||
await self.memory_manager.initialize()
|
||
|
||
# 初始化感知记忆层
|
||
self.perceptual_manager = PerceptualMemoryManager(
|
||
data_dir=self.data_dir,
|
||
**self._config["perceptual"],
|
||
)
|
||
await self.perceptual_manager.initialize()
|
||
|
||
# 初始化短期记忆层
|
||
self.short_term_manager = ShortTermMemoryManager(
|
||
data_dir=self.data_dir,
|
||
**self._config["short_term"],
|
||
)
|
||
await self.short_term_manager.initialize()
|
||
|
||
# 初始化长期记忆层
|
||
self.long_term_manager = LongTermMemoryManager(
|
||
memory_manager=self.memory_manager,
|
||
**self._config["long_term"],
|
||
)
|
||
await self.long_term_manager.initialize()
|
||
|
||
self._initialized = True
|
||
logger.info("统一记忆管理器初始化完成")
|
||
|
||
# 启动自动转移任务
|
||
self._start_auto_transfer_task()
|
||
|
||
except Exception as e:
|
||
logger.error(f"统一记忆管理器初始化失败: {e}")
|
||
raise
|
||
|
||
async def add_message(self, message: dict[str, Any]) -> MemoryBlock | None:
|
||
"""
|
||
添加消息到感知记忆层
|
||
|
||
Args:
|
||
message: 消息字典
|
||
|
||
Returns:
|
||
如果创建了新块，返回 MemoryBlock
|
||
"""
|
||
if not self._initialized:
|
||
await self.initialize()
|
||
|
||
new_block = await self.perceptual_manager.add_message(message)
|
||
|
||
# 注意：感知→短期的转移由召回触发，不是由添加消息触发
|
||
# 转移逻辑在 search_memories 中处理
|
||
|
||
return new_block
|
||
|
||
# 已移除 _process_activated_blocks 方法
|
||
# 转移逻辑现在在 search_memories 中处理：
|
||
# 当召回某个记忆块时，如果其 recall_count >= activation_threshold，
|
||
# 立即将该块转移到短期记忆
|
||
|
||
async def search_memories(
|
||
self, query_text: str, use_judge: bool = True, recent_chat_history: str = ""
|
||
) -> dict[str, Any]:
|
||
"""
|
||
智能检索记忆
|
||
|
||
流程：
|
||
1. 优先检索感知记忆和短期记忆
|
||
2. 使用裁判模型评估是否充足
|
||
3. 如果不充足，生成补充 query 并检索长期记忆
|
||
|
||
Args:
|
||
query_text: 查询文本
|
||
use_judge: 是否使用裁判模型
|
||
recent_chat_history: 最近的聊天历史上下文（可选）
|
||
|
||
Returns:
|
||
检索结果字典，包含：
|
||
- perceptual_blocks: 感知记忆块列表
|
||
- short_term_memories: 短期记忆列表
|
||
- long_term_memories: 长期记忆列表
|
||
- judge_decision: 裁判决策（如果使用）
|
||
"""
|
||
if not self._initialized:
|
||
await self.initialize()
|
||
|
||
try:
|
||
result = {
|
||
"perceptual_blocks": [],
|
||
"short_term_memories": [],
|
||
"long_term_memories": [],
|
||
"judge_decision": None,
|
||
}
|
||
|
||
# 步骤1: 并行检索感知记忆和短期记忆（优化：消除任务创建开销）
|
||
perceptual_blocks, short_term_memories = await asyncio.gather(
|
||
self.perceptual_manager.recall_blocks(query_text),
|
||
self.short_term_manager.search_memories(query_text),
|
||
)
|
||
|
||
# 步骤1.5: 检查需要转移的感知块，推迟到后台处理（优化：单遍扫描与转移）
|
||
blocks_to_transfer = []
|
||
for block in perceptual_blocks:
|
||
if block.metadata.get("needs_transfer", False):
|
||
block.metadata["needs_transfer"] = False # 立即标记，避免重复
|
||
blocks_to_transfer.append(block)
|
||
|
||
if blocks_to_transfer:
|
||
logger.debug(
|
||
f"检测到 {len(blocks_to_transfer)} 个感知记忆需要转移，已交由后台后处理任务执行"
|
||
)
|
||
self._schedule_perceptual_block_transfer(blocks_to_transfer)
|
||
|
||
result["perceptual_blocks"] = perceptual_blocks
|
||
result["short_term_memories"] = short_term_memories
|
||
|
||
# 步骤2: 裁判模型评估
|
||
if use_judge:
|
||
judge_decision = await self._judge_retrieval_sufficiency(
|
||
query_text, perceptual_blocks, short_term_memories, recent_chat_history
|
||
)
|
||
result["judge_decision"] = judge_decision
|
||
|
||
# 步骤3: 如果不充足，检索长期记忆
|
||
if not judge_decision.is_sufficient:
|
||
logger.info("判官判断记忆不足，开始检索长期记忆")
|
||
|
||
queries = [query_text, *judge_decision.additional_queries]
|
||
long_term_memories = await self._retrieve_long_term_memories(
|
||
base_query=query_text,
|
||
queries=queries,
|
||
recent_chat_history=recent_chat_history,
|
||
)
|
||
|
||
result["long_term_memories"] = long_term_memories
|
||
|
||
else:
|
||
# 不使用裁判，直接检索长期记忆
|
||
long_term_memories = await self.memory_manager.search_memories(
|
||
query=query_text,
|
||
top_k=5,
|
||
use_multi_query=False,
|
||
)
|
||
result["long_term_memories"] = long_term_memories
|
||
|
||
return result
|
||
|
||
except Exception as e:
|
||
logger.error(f"智能检索失败: {e}")
|
||
return {
|
||
"perceptual_blocks": [],
|
||
"short_term_memories": [],
|
||
"long_term_memories": [],
|
||
"error": str(e),
|
||
}
|
||
|
||
async def _judge_retrieval_sufficiency(
|
||
self,
|
||
query: str,
|
||
perceptual_blocks: list[MemoryBlock],
|
||
short_term_memories: list[ShortTermMemory],
|
||
recent_chat_history: str = "",
|
||
) -> JudgeDecision:
|
||
"""
|
||
使用裁判模型评估检索结果是否充足
|
||
|
||
Args:
|
||
query: 原始查询
|
||
perceptual_blocks: 感知记忆块
|
||
short_term_memories: 短期记忆
|
||
recent_chat_history: 最近的聊天历史上下文（可选）
|
||
|
||
Returns:
|
||
裁判决策
|
||
"""
|
||
try:
|
||
from src.config.config import model_config
|
||
from src.llm_models.utils_model import LLMRequest
|
||
from src.memory_graph.utils.three_tier_formatter import memory_formatter
|
||
|
||
# 使用新的三级记忆格式化器
|
||
perceptual_desc = await memory_formatter.format_perceptual_memory(perceptual_blocks)
|
||
short_term_desc = memory_formatter.format_short_term_memory(short_term_memories)
|
||
|
||
# 构建聊天历史块（如果提供）
|
||
chat_history_block = ""
|
||
if recent_chat_history:
|
||
chat_history_block = f"""**最近的聊天历史：**
|
||
{recent_chat_history}
|
||
|
||
"""
|
||
|
||
prompt = f"""你是“记忆判官”（记忆检索评估专家）。你的任务是：基于给定的历史消息、当前消息，以及我们已经检索到的“感知记忆”和“短期记忆”，判断是否还需要检索“长期记忆”（LTM）来支撑一次准确、完整、上下文一致的回复。
|
||
|
||
**总体偏好（重要）：**
|
||
- 我们宁可多花一点资源去检索长期记忆，也不要在本该检索时漏检索。
|
||
- 因此：只要存在明显不确定、信息缺口、或需要更精确细节的情况，就倾向于判定“现有记忆不充足”（`is_sufficient: false`）。
|
||
|
||
**输入：**
|
||
**当前用户消息：**
|
||
{query}
|
||
|
||
{chat_history_block}**已检索到的感知记忆（即时对话，格式：【时间 (聊天流)】消息列表）：**
|
||
{perceptual_desc or '（无）'}
|
||
|
||
**已检索到的短期记忆（结构化信息，自然语言描述）：**
|
||
{short_term_desc or '（无）'}
|
||
|
||
**什么时候必须检索长期记忆（满足任一条 → `is_sufficient: false`）：**
|
||
1. **用户明确要求回忆/找回过去信息**：例如“你还记得…？”“上次我们说到…？”“帮我回忆一下…/之前…/那天…/某次…”
|
||
2. **你对答案没有把握或存在不确定性**：例如无法确定人物/事件/时间/地点/偏好/承诺/任务细节，或只能给出模糊猜测。
|
||
3. **现有记忆不足以给出精确回答**：要给出具体结论、细节、步骤、承诺、时间线、决定依据，但感知/短期记忆缺少关键事实。
|
||
4. **对话依赖用户个体历史**：涉及用户的个人信息、偏好、长期目标、过往经历、已约定事项、持续进行的项目/任务，需要更早的上下文才能回答。
|
||
5. **指代不清或背景缺失**：出现“那个/那件事/他/她/它/之前说的/你知道的”等省略指代，现有记忆不足以唯一指向。
|
||
6. **记忆冲突或碎片化**：感知/短期记忆之间存在矛盾、时间线断裂、或信息片段无法拼成完整图景。
|
||
|
||
**什么时候可以不检索（同时满足全部条件 → `is_sufficient: true`）：**
|
||
- 用户只是闲聊/打招呼/情绪表达/泛化问题（不依赖用户个人历史），且现有记忆已足以给出可靠且一致的回复；
|
||
- 你能在不猜测的情况下回答，且不需要更早的细节来保证准确性。
|
||
|
||
**输出要求（JSON）：**
|
||
- `is_sufficient`: `true` 表示“无需检索长期记忆”；`false` 表示“需要检索长期记忆”
|
||
- `confidence`: 0~1，表示你对该判断的把握；若你偏向检索但仍不确定，也应输出较低/中等置信度并保持 `is_sufficient: false`
|
||
- `missing_aspects`: 列出阻碍精确回答的缺失点（可为空数组）
|
||
- `additional_queries`: 给出 1~5 条用于检索长期记忆的补充 query（尽量短、可检索、包含关键实体/事件/时间线线索；可为空数组）
|
||
|
||
请仅输出 JSON（可以用 ```json 包裹，也可以直接输出纯 JSON）："""
|
||
|
||
# 调用记忆裁判模型
|
||
if not model_config.model_task_config:
|
||
raise ValueError("模型任务配置未加载")
|
||
llm = LLMRequest(
|
||
model_set=model_config.model_task_config.memory_judge,
|
||
request_type="unified_memory.judge",
|
||
)
|
||
|
||
response, _ = await llm.generate_response_async(
|
||
prompt,
|
||
temperature=0.1,
|
||
max_tokens=600,
|
||
)
|
||
|
||
# 解析响应
|
||
import json
|
||
import re
|
||
|
||
json_match = re.search(r"```json\s*(.*?)\s*```", response, re.DOTALL)
|
||
if json_match:
|
||
json_str = json_match.group(1)
|
||
else:
|
||
json_str = response.strip()
|
||
|
||
data = json.loads(json_str)
|
||
|
||
decision = JudgeDecision(
|
||
is_sufficient=data.get("is_sufficient", False),
|
||
confidence=data.get("confidence", 0.5),
|
||
reasoning=data.get("reasoning", ""),
|
||
additional_queries=data.get("additional_queries", []),
|
||
missing_aspects=data.get("missing_aspects", []),
|
||
)
|
||
|
||
return decision
|
||
|
||
except Exception as e:
|
||
logger.error(f"裁判模型评估失败: {e}")
|
||
# 默认判定为不充足，需要检索长期记忆
|
||
return JudgeDecision(
|
||
is_sufficient=False,
|
||
confidence=0.3,
|
||
reasoning=f"裁判模型失败: {e}",
|
||
additional_queries=[query],
|
||
)
|
||
|
||
def _schedule_perceptual_block_transfer(self, blocks: list[MemoryBlock]) -> None:
|
||
"""将感知记忆块转移到短期记忆，后台执行以避免阻塞（优化：避免不必要的列表复制）"""
|
||
if not blocks:
|
||
return
|
||
|
||
# 优化：直接传递 blocks 而不再 list(blocks)
|
||
task = asyncio.create_task(
|
||
self._transfer_blocks_to_short_term(blocks)
|
||
)
|
||
self._attach_background_task_callback(task, "perceptual->short-term transfer")
|
||
|
||
def _attach_background_task_callback(self, task: asyncio.Task, task_name: str) -> None:
|
||
"""确保后台任务异常被记录"""
|
||
|
||
def _callback(done_task: asyncio.Task) -> None:
|
||
try:
|
||
done_task.result()
|
||
except asyncio.CancelledError:
|
||
logger.info(f"{task_name} 后台任务已取消")
|
||
except Exception as exc:
|
||
logger.error(f"{task_name} 后台任务失败: {exc}")
|
||
|
||
task.add_done_callback(_callback)
|
||
|
||
def _trigger_transfer_wakeup(self) -> None:
|
||
"""通知自动转移任务立即检查缓存"""
|
||
if self._transfer_wakeup_event and not self._transfer_wakeup_event.is_set():
|
||
self._transfer_wakeup_event.set()
|
||
|
||
def _calculate_auto_sleep_interval(self) -> float:
|
||
"""根据短期内存压力计算自适应等待间隔（优化：查表法替代链式比较）"""
|
||
base_interval = self._auto_transfer_interval
|
||
if not getattr(self, "short_term_manager", None):
|
||
return base_interval
|
||
|
||
max_memories = max(1, getattr(self.short_term_manager, "max_memories", 1))
|
||
occupancy = len(self.short_term_manager.memories) / max_memories
|
||
|
||
# 优化：使用查表法替代链式 if 判断（O(1) vs O(n)）
|
||
occupancy_thresholds = [
|
||
(0.8, 2.0, 0.1),
|
||
(0.5, 5.0, 0.2),
|
||
(0.3, 10.0, 0.4),
|
||
(0.1, 15.0, 0.6),
|
||
]
|
||
|
||
for threshold, min_val, factor in occupancy_thresholds:
|
||
if occupancy >= threshold:
|
||
return max(min_val, base_interval * factor)
|
||
|
||
return base_interval
|
||
|
||
async def _transfer_blocks_to_short_term(self, blocks: list[MemoryBlock]) -> None:
|
||
"""实际转换逻辑在后台执行（优化：并行处理多个块，批量触发唤醒）"""
|
||
logger.debug(f"正在后台处理 {len(blocks)} 个感知记忆块")
|
||
|
||
# 优化：使用 asyncio.gather 并行处理转移
|
||
async def _transfer_single(block: MemoryBlock) -> tuple[MemoryBlock, bool]:
|
||
try:
|
||
stm = await self.short_term_manager.add_from_block(block)
|
||
if not stm:
|
||
return block, False
|
||
|
||
await self.perceptual_manager.remove_block(block.id)
|
||
logger.debug(f"✓ 记忆块 {block.id} 已被转移到短期记忆 {stm.id}")
|
||
return block, True
|
||
except Exception as exc:
|
||
logger.error(f"后台转移失败，记忆块 {block.id}: {exc}")
|
||
return block, False
|
||
|
||
# 并行处理所有块
|
||
results = await asyncio.gather(*[_transfer_single(block) for block in blocks], return_exceptions=True)
|
||
|
||
# 统计成功的转移
|
||
success_count = sum(1 for result in results if isinstance(result, tuple) and result[1])
|
||
if success_count > 0:
|
||
self._trigger_transfer_wakeup()
|
||
logger.debug(f"✅ 后台转移: 成功 {success_count}/{len(blocks)} 个块")
|
||
|
||
def _build_manual_multi_queries(self, queries: list[str]) -> list[dict[str, float]]:
|
||
"""去重裁判查询并附加权重以进行多查询搜索（优化：使用字典推导式）"""
|
||
# 优化：单遍去重（避免多次 strip 和 in 检查）
|
||
seen = set()
|
||
decay = 0.15
|
||
manual_queries: list[dict[str, Any]] = []
|
||
|
||
for raw in queries:
|
||
text = (raw or "").strip()
|
||
if text and text not in seen:
|
||
seen.add(text)
|
||
weight = max(0.3, 1.0 - len(manual_queries) * decay)
|
||
manual_queries.append({"text": text, "weight": round(weight, 2)})
|
||
|
||
# 过滤单条或空列表
|
||
return manual_queries if len(manual_queries) > 1 else []
|
||
|
||
async def _retrieve_long_term_memories(
|
||
self,
|
||
base_query: str,
|
||
queries: list[str],
|
||
recent_chat_history: str = "",
|
||
) -> list[Any]:
|
||
"""可一次性运行多查询搜索的集中式长期检索条目（优化：减少中间对象创建）"""
|
||
manual_queries = self._build_manual_multi_queries(queries)
|
||
|
||
# 优化：仅在必要时创建 context 字典
|
||
search_params: dict[str, Any] = {
|
||
"query": base_query,
|
||
"top_k": self._config["long_term"]["search_top_k"],
|
||
"use_multi_query": bool(manual_queries),
|
||
}
|
||
|
||
if recent_chat_history or manual_queries:
|
||
context: dict[str, Any] = {}
|
||
if recent_chat_history:
|
||
context["chat_history"] = recent_chat_history
|
||
if manual_queries:
|
||
context["manual_multi_queries"] = manual_queries
|
||
search_params["context"] = context
|
||
|
||
memories = await self.memory_manager.search_memories(**search_params)
|
||
return self._deduplicate_memories(memories)
|
||
|
||
def _deduplicate_memories(self, memories: list[Any]) -> list[Any]:
|
||
"""通过 memory.id 去重（优化：支持 dict 和 object，单遍处理）"""
|
||
seen_ids: set[str] = set()
|
||
unique_memories: list[Any] = []
|
||
|
||
for mem in memories:
|
||
# 支持两种 ID 访问方式
|
||
mem_id = None
|
||
if isinstance(mem, dict):
|
||
mem_id = mem.get("id")
|
||
else:
|
||
mem_id = getattr(mem, "id", None)
|
||
|
||
# 检查去重
|
||
if mem_id and mem_id in seen_ids:
|
||
continue
|
||
|
||
unique_memories.append(mem)
|
||
if mem_id:
|
||
seen_ids.add(mem_id)
|
||
|
||
return unique_memories
|
||
|
||
|
||
def _start_auto_transfer_task(self) -> None:
|
||
"""启动自动转移任务"""
|
||
if self._auto_transfer_task and not self._auto_transfer_task.done():
|
||
logger.warning("自动转移任务已在运行")
|
||
return
|
||
|
||
if self._transfer_wakeup_event is None:
|
||
self._transfer_wakeup_event = asyncio.Event()
|
||
else:
|
||
self._transfer_wakeup_event.clear()
|
||
|
||
self._auto_transfer_task = asyncio.create_task(self._auto_transfer_loop())
|
||
# 立即触发一次检查，避免启动初期的长时间等待
|
||
self._transfer_wakeup_event.set()
|
||
logger.debug("自动转移任务已启动并触发首次检查")
|
||
|
||
async def _auto_transfer_loop(self) -> None:
|
||
"""自动转移循环（简化版：短期记忆满额时整批转移）"""
|
||
|
||
while True:
|
||
try:
|
||
sleep_interval = self._calculate_auto_sleep_interval()
|
||
if self._transfer_wakeup_event is not None:
|
||
try:
|
||
await asyncio.wait_for(
|
||
self._transfer_wakeup_event.wait(),
|
||
timeout=sleep_interval,
|
||
)
|
||
self._transfer_wakeup_event.clear()
|
||
except asyncio.TimeoutError:
|
||
pass
|
||
else:
|
||
await asyncio.sleep(sleep_interval)
|
||
|
||
# 最简单策略：仅当短期记忆满额时，直接整批转移全部短期记忆；没满则不处理
|
||
max_memories = max(1, getattr(self.short_term_manager, "max_memories", 1))
|
||
if len(self.short_term_manager.memories) < max_memories:
|
||
continue
|
||
|
||
batch = list(self.short_term_manager.memories)
|
||
if not batch:
|
||
continue
|
||
|
||
logger.info(
|
||
f"短期记忆已满({len(batch)}/{max_memories})，开始整批转移到长期记忆"
|
||
)
|
||
result = await self.long_term_manager.transfer_from_short_term(batch)
|
||
|
||
if result.get("transferred_memory_ids"):
|
||
await self.short_term_manager.clear_transferred_memories(
|
||
result["transferred_memory_ids"]
|
||
)
|
||
logger.debug(f"✅ 整批转移完成: {result}")
|
||
|
||
except asyncio.CancelledError:
|
||
logger.debug("自动转移循环被取消")
|
||
break
|
||
except Exception as e:
|
||
logger.error(f"自动转移循环异常: {e}")
|
||
|
||
async def manual_transfer(self) -> dict[str, Any]:
|
||
"""
|
||
手动触发短期记忆到长期记忆的转移
|
||
|
||
Returns:
|
||
转移结果
|
||
"""
|
||
if not self._initialized:
|
||
await self.initialize()
|
||
|
||
try:
|
||
max_memories = max(1, getattr(self.short_term_manager, "max_memories", 1))
|
||
if len(self.short_term_manager.memories) < max_memories:
|
||
return {
|
||
"message": f"短期记忆未满({len(self.short_term_manager.memories)}/{max_memories})，不触发转移",
|
||
"transferred_count": 0,
|
||
}
|
||
|
||
memories_to_transfer = list(self.short_term_manager.memories)
|
||
if not memories_to_transfer:
|
||
return {"message": "短期记忆为空，无需转移", "transferred_count": 0}
|
||
|
||
# 执行转移
|
||
result = await self.long_term_manager.transfer_from_short_term(memories_to_transfer)
|
||
|
||
# 清除已转移的记忆
|
||
if result.get("transferred_memory_ids"):
|
||
await self.short_term_manager.clear_transferred_memories(
|
||
result["transferred_memory_ids"]
|
||
)
|
||
|
||
logger.info(f"手动转移完成: {result}")
|
||
return result
|
||
|
||
except Exception as e:
|
||
logger.error(f"手动转移失败: {e}")
|
||
return {"error": str(e), "transferred_count": 0}
|
||
|
||
def get_statistics(self) -> dict[str, Any]:
|
||
"""获取三层记忆系统的统计信息"""
|
||
if not self._initialized:
|
||
return {}
|
||
|
||
return {
|
||
"perceptual": self.perceptual_manager.get_statistics(),
|
||
"short_term": self.short_term_manager.get_statistics(),
|
||
"long_term": self.long_term_manager.get_statistics(),
|
||
"total_system_memories": (
|
||
self.perceptual_manager.get_statistics().get("total_messages", 0)
|
||
+ self.short_term_manager.get_statistics().get("total_memories", 0)
|
||
+ self.long_term_manager.get_statistics().get("total_memories", 0)
|
||
),
|
||
}
|
||
|
||
async def shutdown(self) -> None:
|
||
"""关闭统一记忆管理器"""
|
||
if not self._initialized:
|
||
return
|
||
|
||
try:
|
||
logger.info("正在关闭统一记忆管理器...")
|
||
|
||
# 取消自动转移任务
|
||
if self._auto_transfer_task and not self._auto_transfer_task.done():
|
||
self._auto_transfer_task.cancel()
|
||
try:
|
||
await self._auto_transfer_task
|
||
except asyncio.CancelledError:
|
||
pass
|
||
|
||
# 关闭各层管理器
|
||
if self.perceptual_manager:
|
||
await self.perceptual_manager.shutdown()
|
||
|
||
if self.short_term_manager:
|
||
await self.short_term_manager.shutdown()
|
||
|
||
if self.long_term_manager:
|
||
await self.long_term_manager.shutdown()
|
||
|
||
if self.memory_manager:
|
||
await self.memory_manager.shutdown()
|
||
|
||
self._initialized = False
|
||
logger.info("统一记忆管理器已关闭")
|
||
|
||
except Exception as e:
|
||
logger.error(f"关闭统一记忆管理器失败: {e}")
|