336 lines
10 KiB
Python
336 lines
10 KiB
Python
"""
|
||
记忆图可视化服务器
|
||
|
||
提供 Web API 用于可视化记忆图数据
|
||
"""
|
||
|
||
import asyncio
|
||
import logging
|
||
|
||
# 添加项目根目录到 Python 路径
|
||
import sys
|
||
from datetime import datetime
|
||
from pathlib import Path
|
||
from typing import Optional
|
||
|
||
from flask import Flask, jsonify, render_template, request
|
||
from flask_cors import CORS
|
||
|
||
project_root = Path(__file__).parent.parent.parent
|
||
sys.path.insert(0, str(project_root))
|
||
|
||
from src.memory_graph.manager import MemoryManager
|
||
|
||
logging.basicConfig(level=logging.INFO)
|
||
logger = logging.getLogger(__name__)
|
||
|
||
app = Flask(__name__)
|
||
CORS(app) # 允许跨域请求
|
||
|
||
# 全局记忆管理器
|
||
memory_manager: Optional[MemoryManager] = None
|
||
|
||
|
||
def init_memory_manager():
|
||
"""初始化记忆管理器"""
|
||
global memory_manager
|
||
if memory_manager is None:
|
||
try:
|
||
memory_manager = MemoryManager()
|
||
# 在新的事件循环中初始化
|
||
loop = asyncio.new_event_loop()
|
||
asyncio.set_event_loop(loop)
|
||
loop.run_until_complete(memory_manager.initialize())
|
||
logger.info("记忆管理器初始化成功")
|
||
except Exception as e:
|
||
logger.error(f"初始化记忆管理器失败: {e}")
|
||
raise
|
||
|
||
|
||
@app.route("/")
|
||
def index():
|
||
"""主页面"""
|
||
return render_template("visualizer.html")
|
||
|
||
|
||
@app.route("/api/graph/full")
|
||
def get_full_graph():
|
||
"""
|
||
获取完整记忆图数据
|
||
|
||
返回所有节点和边，格式化为前端可用的结构
|
||
"""
|
||
try:
|
||
if memory_manager is None:
|
||
init_memory_manager()
|
||
|
||
# 获取所有记忆
|
||
loop = asyncio.new_event_loop()
|
||
asyncio.set_event_loop(loop)
|
||
|
||
# 获取所有记忆
|
||
all_memories = memory_manager.graph_store.get_all_memories()
|
||
|
||
# 构建节点和边数据
|
||
nodes_dict = {} # {node_id: node_data}
|
||
edges_dict = {} # {edge_id: edge_data} - 使用字典去重
|
||
memory_info = []
|
||
|
||
for memory in all_memories:
|
||
# 添加记忆信息
|
||
memory_info.append(
|
||
{
|
||
"id": memory.id,
|
||
"type": memory.memory_type.value,
|
||
"importance": memory.importance,
|
||
"activation": memory.activation,
|
||
"status": memory.status.value,
|
||
"created_at": memory.created_at.isoformat(),
|
||
"text": memory.to_text(),
|
||
"access_count": memory.access_count,
|
||
}
|
||
)
|
||
|
||
# 处理节点
|
||
for node in memory.nodes:
|
||
if node.id not in nodes_dict:
|
||
nodes_dict[node.id] = {
|
||
"id": node.id,
|
||
"label": node.content,
|
||
"type": node.node_type.value,
|
||
"group": node.node_type.name, # 用于颜色分组
|
||
"title": f"{node.node_type.value}: {node.content}",
|
||
"metadata": node.metadata,
|
||
"created_at": node.created_at.isoformat(),
|
||
}
|
||
|
||
# 处理边 - 使用字典自动去重
|
||
for edge in memory.edges:
|
||
edge_id = edge.id
|
||
# 如果ID已存在，生成唯一ID
|
||
counter = 1
|
||
original_edge_id = edge_id
|
||
while edge_id in edges_dict:
|
||
edge_id = f"{original_edge_id}_{counter}"
|
||
counter += 1
|
||
|
||
edges_dict[edge_id] = {
|
||
"id": edge_id,
|
||
"from": edge.source_id,
|
||
"to": edge.target_id,
|
||
"label": edge.relation,
|
||
"type": edge.edge_type.value,
|
||
"importance": edge.importance,
|
||
"title": f"{edge.edge_type.value}: {edge.relation}",
|
||
"arrows": "to",
|
||
"memory_id": memory.id,
|
||
}
|
||
|
||
nodes_list = list(nodes_dict.values())
|
||
edges_list = list(edges_dict.values())
|
||
|
||
return jsonify(
|
||
{
|
||
"success": True,
|
||
"data": {
|
||
"nodes": nodes_list,
|
||
"edges": edges_list,
|
||
"memories": memory_info,
|
||
"stats": {
|
||
"total_nodes": len(nodes_list),
|
||
"total_edges": len(edges_list),
|
||
"total_memories": len(all_memories),
|
||
},
|
||
},
|
||
}
|
||
)
|
||
|
||
except Exception as e:
|
||
logger.error(f"获取图数据失败: {e}", exc_info=True)
|
||
return jsonify({"success": False, "error": str(e)}), 500
|
||
|
||
|
||
@app.route("/api/memory/<memory_id>")
|
||
def get_memory_detail(memory_id: str):
|
||
"""
|
||
获取特定记忆的详细信息
|
||
|
||
Args:
|
||
memory_id: 记忆ID
|
||
"""
|
||
try:
|
||
if memory_manager is None:
|
||
init_memory_manager()
|
||
|
||
memory = memory_manager.graph_store.get_memory_by_id(memory_id)
|
||
|
||
if memory is None:
|
||
return jsonify({"success": False, "error": "记忆不存在"}), 404
|
||
|
||
return jsonify({"success": True, "data": memory.to_dict()})
|
||
|
||
except Exception as e:
|
||
logger.error(f"获取记忆详情失败: {e}", exc_info=True)
|
||
return jsonify({"success": False, "error": str(e)}), 500
|
||
|
||
|
||
@app.route("/api/search")
|
||
def search_memories():
|
||
"""
|
||
搜索记忆
|
||
|
||
Query参数:
|
||
- q: 搜索关键词
|
||
- type: 记忆类型过滤
|
||
- limit: 返回数量限制
|
||
"""
|
||
try:
|
||
if memory_manager is None:
|
||
init_memory_manager()
|
||
|
||
query = request.args.get("q", "")
|
||
memory_type = request.args.get("type", None)
|
||
limit = int(request.args.get("limit", 50))
|
||
|
||
loop = asyncio.new_event_loop()
|
||
asyncio.set_event_loop(loop)
|
||
|
||
# 执行搜索
|
||
results = loop.run_until_complete(memory_manager.search_memories(query=query, top_k=limit))
|
||
|
||
# 构建返回数据
|
||
memories = []
|
||
for memory in results:
|
||
memories.append(
|
||
{
|
||
"id": memory.id,
|
||
"text": memory.to_text(),
|
||
"type": memory.memory_type.value,
|
||
"importance": memory.importance,
|
||
"created_at": memory.created_at.isoformat(),
|
||
}
|
||
)
|
||
|
||
return jsonify(
|
||
{
|
||
"success": True,
|
||
"data": {
|
||
"results": memories,
|
||
"count": len(memories),
|
||
},
|
||
}
|
||
)
|
||
|
||
except Exception as e:
|
||
logger.error(f"搜索失败: {e}", exc_info=True)
|
||
return jsonify({"success": False, "error": str(e)}), 500
|
||
|
||
|
||
@app.route("/api/stats")
|
||
def get_statistics():
|
||
"""
|
||
获取记忆图统计信息
|
||
"""
|
||
try:
|
||
if memory_manager is None:
|
||
init_memory_manager()
|
||
|
||
# 获取统计信息
|
||
all_memories = memory_manager.graph_store.get_all_memories()
|
||
all_nodes = set()
|
||
all_edges = 0
|
||
|
||
for memory in all_memories:
|
||
for node in memory.nodes:
|
||
all_nodes.add(node.id)
|
||
all_edges += len(memory.edges)
|
||
|
||
stats = {
|
||
"total_memories": len(all_memories),
|
||
"total_nodes": len(all_nodes),
|
||
"total_edges": all_edges,
|
||
"node_types": {},
|
||
"memory_types": {},
|
||
}
|
||
|
||
# 统计节点类型分布
|
||
for memory in all_memories:
|
||
mem_type = memory.memory_type.value
|
||
stats["memory_types"][mem_type] = stats["memory_types"].get(mem_type, 0) + 1
|
||
|
||
for node in memory.nodes:
|
||
node_type = node.node_type.value
|
||
stats["node_types"][node_type] = stats["node_types"].get(node_type, 0) + 1
|
||
|
||
return jsonify({"success": True, "data": stats})
|
||
|
||
except Exception as e:
|
||
logger.error(f"获取统计信息失败: {e}", exc_info=True)
|
||
return jsonify({"success": False, "error": str(e)}), 500
|
||
|
||
|
||
@app.route("/api/files")
|
||
def list_files():
|
||
"""
|
||
列出所有可用的数据文件
|
||
注意: 完整版服务器直接使用内存中的数据，不支持文件切换
|
||
"""
|
||
try:
|
||
from pathlib import Path
|
||
|
||
data_dir = Path("data/memory_graph")
|
||
|
||
files = []
|
||
if data_dir.exists():
|
||
for f in data_dir.glob("*.json"):
|
||
stat = f.stat()
|
||
files.append(
|
||
{
|
||
"path": str(f),
|
||
"name": f.name,
|
||
"size": stat.st_size,
|
||
"size_kb": round(stat.st_size / 1024, 2),
|
||
"modified": datetime.fromtimestamp(stat.st_mtime).isoformat(),
|
||
"modified_readable": datetime.fromtimestamp(stat.st_mtime).strftime("%Y-%m-%d %H:%M:%S"),
|
||
"is_current": True, # 完整版始终使用内存数据
|
||
}
|
||
)
|
||
|
||
return jsonify(
|
||
{
|
||
"success": True,
|
||
"files": files,
|
||
"count": len(files),
|
||
"current_file": "memory_manager (实时数据)",
|
||
"note": "完整版服务器使用实时内存数据，如需切换文件请使用独立版服务器",
|
||
}
|
||
)
|
||
except Exception as e:
|
||
logger.error(f"获取文件列表失败: {e}", exc_info=True)
|
||
return jsonify({"success": False, "error": str(e)}), 500
|
||
|
||
|
||
@app.route("/api/reload")
|
||
def reload_data():
|
||
"""
|
||
重新加载数据
|
||
"""
|
||
return jsonify({"success": True, "message": "完整版服务器使用实时数据，无需重新加载", "note": "数据始终是最新的"})
|
||
|
||
|
||
def run_server(host: str = "127.0.0.1", port: int = 5000, debug: bool = False):
|
||
"""
|
||
启动可视化服务器
|
||
|
||
Args:
|
||
host: 服务器地址
|
||
port: 端口号
|
||
debug: 是否开启调试模式
|
||
"""
|
||
logger.info(f"启动记忆图可视化服务器: http://{host}:{port}")
|
||
app.run(host=host, port=port, debug=debug)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
run_server(debug=True)
|