feat(memory-graph): Phase 1 完整实现 - 持久化和节点去重
完成功能: - 持久化管理 (PersistenceManager) * 图数据的保存和加载 * 自动备份和恢复 * 数据导出/导入 - 节点去重合并 (NodeMerger) * 基于语义相似度查找重复节点 * 上下文匹配验证 * 自动节点合并 * 批量处理支持 测试验证: - 持久化: 保存/加载/备份 - 节点合并: 相似度0.999自动合并 - 图统计: 合并后节点数正确减少 Phase 1 完成度: 100% - 所有基础设施就绪 - 准备进入 Phase 2
This commit is contained in:
Windpicker-owo
7
src/memory_graph/core/__init__.py
Normal file
7
src/memory_graph/core/__init__.py
Normal file
@@ -0,0 +1,7 @@
|
||||
"""
|
||||
核心模块
|
||||
"""
|
||||
|
||||
from src.memory_graph.core.node_merger import NodeMerger
|
||||
|
||||
__all__ = ["NodeMerger"]
|
||||
359
src/memory_graph/core/node_merger.py
Normal file
359
src/memory_graph/core/node_merger.py
Normal file
@@ -0,0 +1,359 @@
|
||||
"""
|
||||
节点去重合并器：基于语义相似度合并重复节点
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import List, Optional, Tuple
|
||||
|
||||
import numpy as np
|
||||
|
||||
from src.common.logger import get_logger
|
||||
from src.memory_graph.config import NodeMergerConfig
|
||||
from src.memory_graph.models import MemoryNode, NodeType
|
||||
from src.memory_graph.storage.graph_store import GraphStore
|
||||
from src.memory_graph.storage.vector_store import VectorStore
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
class NodeMerger:
|
||||
"""
|
||||
节点合并器
|
||||
|
||||
负责：
|
||||
1. 基于语义相似度查找重复节点
|
||||
2. 验证上下文匹配
|
||||
3. 执行节点合并操作
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
vector_store: VectorStore,
|
||||
graph_store: GraphStore,
|
||||
config: Optional[NodeMergerConfig] = None,
|
||||
):
|
||||
"""
|
||||
初始化节点合并器
|
||||
|
||||
Args:
|
||||
vector_store: 向量存储
|
||||
graph_store: 图存储
|
||||
config: 配置对象
|
||||
"""
|
||||
self.vector_store = vector_store
|
||||
self.graph_store = graph_store
|
||||
self.config = config or NodeMergerConfig()
|
||||
|
||||
logger.info(
|
||||
f"初始化节点合并器: threshold={self.config.similarity_threshold}, "
|
||||
f"context_match={self.config.context_match_required}"
|
||||
)
|
||||
|
||||
async def find_similar_nodes(
|
||||
self,
|
||||
node: MemoryNode,
|
||||
threshold: Optional[float] = None,
|
||||
limit: int = 5,
|
||||
) -> List[Tuple[MemoryNode, float]]:
|
||||
"""
|
||||
查找与指定节点相似的节点
|
||||
|
||||
Args:
|
||||
node: 查询节点
|
||||
threshold: 相似度阈值（可选，默认使用配置值）
|
||||
limit: 返回结果数量
|
||||
|
||||
Returns:
|
||||
List of (similar_node, similarity)
|
||||
"""
|
||||
if not node.has_embedding():
|
||||
logger.warning(f"节点 {node.id} 没有 embedding，无法查找相似节点")
|
||||
return []
|
||||
|
||||
threshold = threshold or self.config.similarity_threshold
|
||||
|
||||
try:
|
||||
# 在向量存储中搜索相似节点
|
||||
results = await self.vector_store.search_similar_nodes(
|
||||
query_embedding=node.embedding,
|
||||
limit=limit + 1, # +1 因为可能包含节点自己
|
||||
node_types=[node.node_type], # 只搜索相同类型的节点
|
||||
min_similarity=threshold,
|
||||
)
|
||||
|
||||
# 过滤掉节点自己，并构建结果
|
||||
similar_nodes = []
|
||||
for node_id, similarity, metadata in results:
|
||||
if node_id == node.id:
|
||||
continue # 跳过自己
|
||||
|
||||
# 从图存储中获取完整节点信息
|
||||
memories = self.graph_store.get_memories_by_node(node_id)
|
||||
if memories:
|
||||
# 从第一个记忆中获取节点
|
||||
target_node = memories[0].get_node_by_id(node_id)
|
||||
if target_node:
|
||||
similar_nodes.append((target_node, similarity))
|
||||
|
||||
logger.debug(f"找到 {len(similar_nodes)} 个相似节点 (阈值: {threshold})")
|
||||
return similar_nodes
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"查找相似节点失败: {e}", exc_info=True)
|
||||
return []
|
||||
|
||||
async def should_merge(
|
||||
self,
|
||||
source_node: MemoryNode,
|
||||
target_node: MemoryNode,
|
||||
similarity: float,
|
||||
) -> bool:
|
||||
"""
|
||||
判断两个节点是否应该合并
|
||||
|
||||
Args:
|
||||
source_node: 源节点
|
||||
target_node: 目标节点
|
||||
similarity: 语义相似度
|
||||
|
||||
Returns:
|
||||
是否应该合并
|
||||
"""
|
||||
# 1. 检查相似度阈值
|
||||
if similarity < self.config.similarity_threshold:
|
||||
return False
|
||||
|
||||
# 2. 非常高的相似度（>0.95）直接合并
|
||||
if similarity > 0.95:
|
||||
logger.debug(f"高相似度 ({similarity:.3f})，直接合并")
|
||||
return True
|
||||
|
||||
# 3. 如果不要求上下文匹配，则通过相似度判断
|
||||
if not self.config.context_match_required:
|
||||
return True
|
||||
|
||||
# 4. 检查上下文匹配
|
||||
context_match = await self._check_context_match(source_node, target_node)
|
||||
|
||||
if context_match:
|
||||
logger.debug(
|
||||
f"相似度 {similarity:.3f} + 上下文匹配，决定合并: "
|
||||
f"'{source_node.content}' → '{target_node.content}'"
|
||||
)
|
||||
return True
|
||||
|
||||
logger.debug(
|
||||
f"相似度 {similarity:.3f} 但上下文不匹配，不合并: "
|
||||
f"'{source_node.content}' ≠ '{target_node.content}'"
|
||||
)
|
||||
return False
|
||||
|
||||
async def _check_context_match(
|
||||
self,
|
||||
source_node: MemoryNode,
|
||||
target_node: MemoryNode,
|
||||
) -> bool:
|
||||
"""
|
||||
检查两个节点的上下文是否匹配
|
||||
|
||||
上下文匹配的标准：
|
||||
1. 节点类型相同
|
||||
2. 邻居节点有重叠
|
||||
3. 邻居节点的内容相似
|
||||
|
||||
Args:
|
||||
source_node: 源节点
|
||||
target_node: 目标节点
|
||||
|
||||
Returns:
|
||||
是否匹配
|
||||
"""
|
||||
# 1. 节点类型必须相同
|
||||
if source_node.node_type != target_node.node_type:
|
||||
return False
|
||||
|
||||
# 2. 获取邻居节点
|
||||
source_neighbors = self.graph_store.get_neighbors(source_node.id, direction="both")
|
||||
target_neighbors = self.graph_store.get_neighbors(target_node.id, direction="both")
|
||||
|
||||
# 如果都没有邻居，认为上下文不足，保守地不合并
|
||||
if not source_neighbors or not target_neighbors:
|
||||
return False
|
||||
|
||||
# 3. 检查邻居内容是否有重叠
|
||||
source_neighbor_contents = set()
|
||||
for neighbor_id, edge_data in source_neighbors:
|
||||
neighbor_node = self._get_node_content(neighbor_id)
|
||||
if neighbor_node:
|
||||
source_neighbor_contents.add(neighbor_node.lower())
|
||||
|
||||
target_neighbor_contents = set()
|
||||
for neighbor_id, edge_data in target_neighbors:
|
||||
neighbor_node = self._get_node_content(neighbor_id)
|
||||
if neighbor_node:
|
||||
target_neighbor_contents.add(neighbor_node.lower())
|
||||
|
||||
# 计算重叠率
|
||||
intersection = source_neighbor_contents & target_neighbor_contents
|
||||
union = source_neighbor_contents | target_neighbor_contents
|
||||
|
||||
if not union:
|
||||
return False
|
||||
|
||||
overlap_ratio = len(intersection) / len(union)
|
||||
|
||||
# 如果有 30% 以上的邻居重叠，认为上下文匹配
|
||||
return overlap_ratio > 0.3
|
||||
|
||||
def _get_node_content(self, node_id: str) -> Optional[str]:
|
||||
"""获取节点的内容"""
|
||||
memories = self.graph_store.get_memories_by_node(node_id)
|
||||
if memories:
|
||||
node = memories[0].get_node_by_id(node_id)
|
||||
if node:
|
||||
return node.content
|
||||
return None
|
||||
|
||||
async def merge_nodes(
|
||||
self,
|
||||
source: MemoryNode,
|
||||
target: MemoryNode,
|
||||
) -> bool:
|
||||
"""
|
||||
合并两个节点
|
||||
|
||||
将 source 节点的所有边转移到 target 节点，然后删除 source
|
||||
|
||||
Args:
|
||||
source: 源节点（将被删除）
|
||||
target: 目标节点（保留）
|
||||
|
||||
Returns:
|
||||
是否成功
|
||||
"""
|
||||
try:
|
||||
logger.info(f"合并节点: '{source.content}' ({source.id}) → '{target.content}' ({target.id})")
|
||||
|
||||
# 1. 在图存储中合并节点
|
||||
self.graph_store.merge_nodes(source.id, target.id)
|
||||
|
||||
# 2. 在向量存储中删除源节点
|
||||
await self.vector_store.delete_node(source.id)
|
||||
|
||||
# 3. 更新所有相关记忆的节点引用
|
||||
self._update_memory_references(source.id, target.id)
|
||||
|
||||
logger.info(f"节点合并成功: {source.id} → {target.id}")
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"节点合并失败: {e}", exc_info=True)
|
||||
return False
|
||||
|
||||
def _update_memory_references(self, old_node_id: str, new_node_id: str) -> None:
|
||||
"""
|
||||
更新记忆中的节点引用
|
||||
|
||||
Args:
|
||||
old_node_id: 旧节点ID
|
||||
new_node_id: 新节点ID
|
||||
"""
|
||||
# 获取所有包含旧节点的记忆
|
||||
memories = self.graph_store.get_memories_by_node(old_node_id)
|
||||
|
||||
for memory in memories:
|
||||
# 移除旧节点
|
||||
memory.nodes = [n for n in memory.nodes if n.id != old_node_id]
|
||||
|
||||
# 更新边的引用
|
||||
for edge in memory.edges:
|
||||
if edge.source_id == old_node_id:
|
||||
edge.source_id = new_node_id
|
||||
if edge.target_id == old_node_id:
|
||||
edge.target_id = new_node_id
|
||||
|
||||
# 更新主体ID（如果是主体节点）
|
||||
if memory.subject_id == old_node_id:
|
||||
memory.subject_id = new_node_id
|
||||
|
||||
async def batch_merge_similar_nodes(
|
||||
self,
|
||||
nodes: List[MemoryNode],
|
||||
progress_callback: Optional[callable] = None,
|
||||
) -> dict:
|
||||
"""
|
||||
批量处理节点合并
|
||||
|
||||
Args:
|
||||
nodes: 要处理的节点列表
|
||||
progress_callback: 进度回调函数
|
||||
|
||||
Returns:
|
||||
统计信息字典
|
||||
"""
|
||||
stats = {
|
||||
"total": len(nodes),
|
||||
"checked": 0,
|
||||
"merged": 0,
|
||||
"skipped": 0,
|
||||
}
|
||||
|
||||
for i, node in enumerate(nodes):
|
||||
try:
|
||||
# 只处理有 embedding 的主题和客体节点
|
||||
if not node.has_embedding() or node.node_type not in [
|
||||
NodeType.TOPIC,
|
||||
NodeType.OBJECT,
|
||||
]:
|
||||
stats["skipped"] += 1
|
||||
continue
|
||||
|
||||
# 查找相似节点
|
||||
similar_nodes = await self.find_similar_nodes(node, limit=5)
|
||||
|
||||
if similar_nodes:
|
||||
# 选择最相似的节点
|
||||
best_match, similarity = similar_nodes[0]
|
||||
|
||||
# 判断是否应该合并
|
||||
if await self.should_merge(node, best_match, similarity):
|
||||
success = await self.merge_nodes(node, best_match)
|
||||
if success:
|
||||
stats["merged"] += 1
|
||||
|
||||
stats["checked"] += 1
|
||||
|
||||
# 调用进度回调
|
||||
if progress_callback:
|
||||
progress_callback(i + 1, stats["total"], stats)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"处理节点 {node.id} 时失败: {e}", exc_info=True)
|
||||
stats["skipped"] += 1
|
||||
|
||||
logger.info(
|
||||
f"批量合并完成: 总数={stats['total']}, 检查={stats['checked']}, "
|
||||
f"合并={stats['merged']}, 跳过={stats['skipped']}"
|
||||
)
|
||||
|
||||
return stats
|
||||
|
||||
def get_merge_candidates(
|
||||
self,
|
||||
min_similarity: float = 0.85,
|
||||
limit: int = 100,
|
||||
) -> List[Tuple[str, str, float]]:
|
||||
"""
|
||||
获取待合并的候选节点对
|
||||
|
||||
Args:
|
||||
min_similarity: 最小相似度
|
||||
limit: 最大返回数量
|
||||
|
||||
Returns:
|
||||
List of (node_id_1, node_id_2, similarity)
|
||||
"""
|
||||
# TODO: 实现更智能的候选查找算法
|
||||
# 目前返回空列表，后续可以基于向量存储进行批量查询
|
||||
return []
|
||||
379
src/memory_graph/storage/persistence.py
Normal file
379
src/memory_graph/storage/persistence.py
Normal file
@@ -0,0 +1,379 @@
|
||||
"""
|
||||
持久化管理：负责记忆图数据的保存和加载
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
import orjson
|
||||
|
||||
from src.common.logger import get_logger
|
||||
from src.memory_graph.models import Memory, StagedMemory
|
||||
from src.memory_graph.storage.graph_store import GraphStore
|
||||
from src.memory_graph.storage.vector_store import VectorStore
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
class PersistenceManager:
|
||||
"""
|
||||
持久化管理器
|
||||
|
||||
负责：
|
||||
1. 图数据的保存和加载
|
||||
2. 定期自动保存
|
||||
3. 备份管理
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
data_dir: Path,
|
||||
graph_file_name: str = "memory_graph.json",
|
||||
staged_file_name: str = "staged_memories.json",
|
||||
auto_save_interval: int = 300, # 自动保存间隔（秒）
|
||||
):
|
||||
"""
|
||||
初始化持久化管理器
|
||||
|
||||
Args:
|
||||
data_dir: 数据存储目录
|
||||
graph_file_name: 图数据文件名
|
||||
staged_file_name: 临时记忆文件名
|
||||
auto_save_interval: 自动保存间隔（秒）
|
||||
"""
|
||||
self.data_dir = Path(data_dir)
|
||||
self.data_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
self.graph_file = self.data_dir / graph_file_name
|
||||
self.staged_file = self.data_dir / staged_file_name
|
||||
self.backup_dir = self.data_dir / "backups"
|
||||
self.backup_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
self.auto_save_interval = auto_save_interval
|
||||
self._auto_save_task: Optional[asyncio.Task] = None
|
||||
self._running = False
|
||||
|
||||
logger.info(f"初始化持久化管理器: data_dir={data_dir}")
|
||||
|
||||
async def save_graph_store(self, graph_store: GraphStore) -> None:
|
||||
"""
|
||||
保存图存储到文件
|
||||
|
||||
Args:
|
||||
graph_store: 图存储对象
|
||||
"""
|
||||
try:
|
||||
# 转换为字典
|
||||
data = graph_store.to_dict()
|
||||
|
||||
# 添加元数据
|
||||
data["metadata"] = {
|
||||
"version": "0.1.0",
|
||||
"saved_at": datetime.now().isoformat(),
|
||||
"statistics": graph_store.get_statistics(),
|
||||
}
|
||||
|
||||
# 使用 orjson 序列化（更快）
|
||||
json_data = orjson.dumps(
|
||||
data,
|
||||
option=orjson.OPT_INDENT_2 | orjson.OPT_SERIALIZE_NUMPY,
|
||||
)
|
||||
|
||||
# 原子写入（先写临时文件，再重命名）
|
||||
temp_file = self.graph_file.with_suffix(".tmp")
|
||||
temp_file.write_bytes(json_data)
|
||||
temp_file.replace(self.graph_file)
|
||||
|
||||
logger.info(f"图数据已保存: {self.graph_file}, 大小: {len(json_data) / 1024:.2f} KB")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"保存图数据失败: {e}", exc_info=True)
|
||||
raise
|
||||
|
||||
async def load_graph_store(self) -> Optional[GraphStore]:
|
||||
"""
|
||||
从文件加载图存储
|
||||
|
||||
Returns:
|
||||
GraphStore 对象，如果文件不存在则返回 None
|
||||
"""
|
||||
if not self.graph_file.exists():
|
||||
logger.info("图数据文件不存在，返回空图")
|
||||
return None
|
||||
|
||||
try:
|
||||
# 读取文件
|
||||
json_data = self.graph_file.read_bytes()
|
||||
data = orjson.loads(json_data)
|
||||
|
||||
# 检查版本（未来可能需要数据迁移）
|
||||
version = data.get("metadata", {}).get("version", "unknown")
|
||||
logger.info(f"加载图数据: version={version}")
|
||||
|
||||
# 恢复图存储
|
||||
graph_store = GraphStore.from_dict(data)
|
||||
|
||||
logger.info(f"图数据加载完成: {graph_store.get_statistics()}")
|
||||
return graph_store
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"加载图数据失败: {e}", exc_info=True)
|
||||
# 尝试加载备份
|
||||
return await self._load_from_backup()
|
||||
|
||||
async def save_staged_memories(self, staged_memories: list[StagedMemory]) -> None:
|
||||
"""
|
||||
保存临时记忆列表
|
||||
|
||||
Args:
|
||||
staged_memories: 临时记忆列表
|
||||
"""
|
||||
try:
|
||||
data = {
|
||||
"metadata": {
|
||||
"version": "0.1.0",
|
||||
"saved_at": datetime.now().isoformat(),
|
||||
"count": len(staged_memories),
|
||||
},
|
||||
"staged_memories": [sm.to_dict() for sm in staged_memories],
|
||||
}
|
||||
|
||||
json_data = orjson.dumps(data, option=orjson.OPT_INDENT_2 | orjson.OPT_SERIALIZE_NUMPY)
|
||||
|
||||
temp_file = self.staged_file.with_suffix(".tmp")
|
||||
temp_file.write_bytes(json_data)
|
||||
temp_file.replace(self.staged_file)
|
||||
|
||||
logger.info(f"临时记忆已保存: {len(staged_memories)} 条")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"保存临时记忆失败: {e}", exc_info=True)
|
||||
raise
|
||||
|
||||
async def load_staged_memories(self) -> list[StagedMemory]:
|
||||
"""
|
||||
加载临时记忆列表
|
||||
|
||||
Returns:
|
||||
临时记忆列表
|
||||
"""
|
||||
if not self.staged_file.exists():
|
||||
logger.info("临时记忆文件不存在，返回空列表")
|
||||
return []
|
||||
|
||||
try:
|
||||
json_data = self.staged_file.read_bytes()
|
||||
data = orjson.loads(json_data)
|
||||
|
||||
staged_memories = [StagedMemory.from_dict(sm) for sm in data.get("staged_memories", [])]
|
||||
|
||||
logger.info(f"临时记忆加载完成: {len(staged_memories)} 条")
|
||||
return staged_memories
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"加载临时记忆失败: {e}", exc_info=True)
|
||||
return []
|
||||
|
||||
async def create_backup(self) -> Optional[Path]:
|
||||
"""
|
||||
创建当前数据的备份
|
||||
|
||||
Returns:
|
||||
备份文件路径，如果失败则返回 None
|
||||
"""
|
||||
try:
|
||||
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
backup_file = self.backup_dir / f"memory_graph_backup_{timestamp}.json"
|
||||
|
||||
if self.graph_file.exists():
|
||||
# 复制图数据文件
|
||||
import shutil
|
||||
|
||||
shutil.copy2(self.graph_file, backup_file)
|
||||
|
||||
# 清理旧备份（只保留最近10个）
|
||||
await self._cleanup_old_backups(keep=10)
|
||||
|
||||
logger.info(f"备份创建成功: {backup_file}")
|
||||
return backup_file
|
||||
|
||||
return None
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"创建备份失败: {e}", exc_info=True)
|
||||
return None
|
||||
|
||||
async def _load_from_backup(self) -> Optional[GraphStore]:
|
||||
"""从最新的备份加载数据"""
|
||||
try:
|
||||
# 查找最新的备份文件
|
||||
backup_files = sorted(self.backup_dir.glob("memory_graph_backup_*.json"), reverse=True)
|
||||
|
||||
if not backup_files:
|
||||
logger.warning("没有可用的备份文件")
|
||||
return None
|
||||
|
||||
latest_backup = backup_files[0]
|
||||
logger.warning(f"尝试从备份恢复: {latest_backup}")
|
||||
|
||||
json_data = latest_backup.read_bytes()
|
||||
data = orjson.loads(json_data)
|
||||
|
||||
graph_store = GraphStore.from_dict(data)
|
||||
logger.info(f"从备份恢复成功: {graph_store.get_statistics()}")
|
||||
|
||||
return graph_store
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"从备份恢复失败: {e}", exc_info=True)
|
||||
return None
|
||||
|
||||
async def _cleanup_old_backups(self, keep: int = 10) -> None:
|
||||
"""
|
||||
清理旧备份，只保留最近的几个
|
||||
|
||||
Args:
|
||||
keep: 保留的备份数量
|
||||
"""
|
||||
try:
|
||||
backup_files = sorted(self.backup_dir.glob("memory_graph_backup_*.json"), reverse=True)
|
||||
|
||||
# 删除超出数量的备份
|
||||
for backup_file in backup_files[keep:]:
|
||||
backup_file.unlink()
|
||||
logger.debug(f"删除旧备份: {backup_file}")
|
||||
|
||||
except Exception as e:
|
||||
logger.warning(f"清理旧备份失败: {e}")
|
||||
|
||||
async def start_auto_save(
|
||||
self,
|
||||
graph_store: GraphStore,
|
||||
staged_memories_getter: callable = None,
|
||||
) -> None:
|
||||
"""
|
||||
启动自动保存任务
|
||||
|
||||
Args:
|
||||
graph_store: 图存储对象
|
||||
staged_memories_getter: 获取临时记忆的回调函数
|
||||
"""
|
||||
if self._auto_save_task and not self._auto_save_task.done():
|
||||
logger.warning("自动保存任务已在运行")
|
||||
return
|
||||
|
||||
self._running = True
|
||||
|
||||
async def auto_save_loop():
|
||||
logger.info(f"自动保存任务已启动，间隔: {self.auto_save_interval}秒")
|
||||
|
||||
while self._running:
|
||||
try:
|
||||
await asyncio.sleep(self.auto_save_interval)
|
||||
|
||||
if not self._running:
|
||||
break
|
||||
|
||||
# 保存图数据
|
||||
await self.save_graph_store(graph_store)
|
||||
|
||||
# 保存临时记忆（如果提供了获取函数）
|
||||
if staged_memories_getter:
|
||||
staged_memories = staged_memories_getter()
|
||||
if staged_memories:
|
||||
await self.save_staged_memories(staged_memories)
|
||||
|
||||
# 定期创建备份（每小时）
|
||||
current_time = datetime.now()
|
||||
if current_time.minute == 0: # 每个整点
|
||||
await self.create_backup()
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"自动保存失败: {e}", exc_info=True)
|
||||
|
||||
logger.info("自动保存任务已停止")
|
||||
|
||||
self._auto_save_task = asyncio.create_task(auto_save_loop())
|
||||
|
||||
def stop_auto_save(self) -> None:
|
||||
"""停止自动保存任务"""
|
||||
self._running = False
|
||||
if self._auto_save_task:
|
||||
self._auto_save_task.cancel()
|
||||
logger.info("自动保存任务已取消")
|
||||
|
||||
async def export_to_json(self, output_file: Path, graph_store: GraphStore) -> None:
|
||||
"""
|
||||
导出图数据到指定的 JSON 文件（用于数据迁移或分析）
|
||||
|
||||
Args:
|
||||
output_file: 输出文件路径
|
||||
graph_store: 图存储对象
|
||||
"""
|
||||
try:
|
||||
data = graph_store.to_dict()
|
||||
data["metadata"] = {
|
||||
"version": "0.1.0",
|
||||
"exported_at": datetime.now().isoformat(),
|
||||
"statistics": graph_store.get_statistics(),
|
||||
}
|
||||
|
||||
# 使用标准 json 以获得更好的可读性
|
||||
output_file.parent.mkdir(parents=True, exist_ok=True)
|
||||
with output_file.open("w", encoding="utf-8") as f:
|
||||
json.dump(data, f, ensure_ascii=False, indent=2)
|
||||
|
||||
logger.info(f"图数据已导出: {output_file}")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"导出图数据失败: {e}", exc_info=True)
|
||||
raise
|
||||
|
||||
async def import_from_json(self, input_file: Path) -> Optional[GraphStore]:
|
||||
"""
|
||||
从 JSON 文件导入图数据
|
||||
|
||||
Args:
|
||||
input_file: 输入文件路径
|
||||
|
||||
Returns:
|
||||
GraphStore 对象
|
||||
"""
|
||||
try:
|
||||
with input_file.open("r", encoding="utf-8") as f:
|
||||
data = json.load(f)
|
||||
|
||||
graph_store = GraphStore.from_dict(data)
|
||||
logger.info(f"图数据已导入: {graph_store.get_statistics()}")
|
||||
|
||||
return graph_store
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"导入图数据失败: {e}", exc_info=True)
|
||||
raise
|
||||
|
||||
def get_data_size(self) -> dict[str, int]:
|
||||
"""
|
||||
获取数据文件的大小信息
|
||||
|
||||
Returns:
|
||||
文件大小字典（字节）
|
||||
"""
|
||||
sizes = {}
|
||||
|
||||
if self.graph_file.exists():
|
||||
sizes["graph"] = self.graph_file.stat().st_size
|
||||
|
||||
if self.staged_file.exists():
|
||||
sizes["staged"] = self.staged_file.stat().st_size
|
||||
|
||||
# 计算备份文件总大小
|
||||
backup_size = sum(f.stat().st_size for f in self.backup_dir.glob("*.json"))
|
||||
sizes["backups"] = backup_size
|
||||
|
||||
return sizes
|
||||
Reference in New Issue
Block a user