You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

156 lines
4.6 KiB

4 months ago
import robyn
from robyn import Robyn, jsonify, Response
from typing import Optional, List, Any, Dict
from util.neo4j_utils import Neo4jUtil
from util.neo4j_utils import neo4j_client
def convert_node_to_g6_v5(neo4j_node: dict) -> dict:
node_id = neo4j_node.get("id")
if node_id is None:
raise ValueError("节点必须包含 'id' 字段")
data = {k: v for k, v in neo4j_node.items() if k != "id"}
if "name" not in data and "label" not in data:
data["name"] = str(node_id)
return {
"id": node_id,
"data": data,
"states": [],
"combo": None
}
def build_g6_graph_data_from_results(
nodes: List[Dict[str, Any]],
relationships: List[Dict[str, Any]]
) -> dict:
"""
通用方法根据节点列表和关系列表构建 G6 v5 图数据
Args:
nodes: 节点列表每个节点需含 "id"
relationships: 关系列表每个关系需含:
- source: {"id": ..., ...}
- target: {"id": ..., ...}
- relationship: {"type": str, "properties": dict} 或直接扁平化字段
Returns:
{"nodes": [...], "edges": [...]}
"""
g6_node_map = {}
# 处理显式传入的节点
for node in nodes:
node_id = node.get("id")
if node_id:
g6_node_map[node_id] = convert_node_to_g6_v5(node)
g6_edges = []
for rel in relationships:
source_node = rel.get("source")
target_node = rel.get("target")
if not source_node or not target_node:
continue
source_id = source_node.get("id")
target_id = target_node.get("id")
if not source_id or not target_id:
continue
# 确保 source/target 节点也加入图中(即使未在 nodes 中显式提供)
if source_id not in g6_node_map:
g6_node_map[source_id] = convert_node_to_g6_v5(source_node)
if target_id not in g6_node_map:
g6_node_map[target_id] = convert_node_to_g6_v5(target_node)
# 构建 edge data
edge_data = {}
rel_type_str = rel.get("type") or rel.get("relationship") # 兼容不同结构
if rel_type_str:
edge_data["relationship"] = rel_type_str
# 尝试从 relProps 或 properties 或顶层提取关系属性
rel_props = (
rel.get("relProps") or
rel.get("properties") or
{k: v for k, v in rel.items() if k not in ("source", "target", "type", "relationship")}
)
if isinstance(rel_props, dict):
edge_data.update(rel_props)
g6_edge = {
"source": source_id,
"target": target_id,
"type": "line",
"data": edge_data,
"states": []
}
g6_edges.append(g6_edge)
return {
"nodes": list(g6_node_map.values()),
"edges": g6_edges
}
def build_g6_subgraph_by_props(
neo4j_util: Neo4jUtil,
node_label: str,
node_properties: Dict[str, Any],
direction: str = "both",
rel_type: Optional[str] = None
) -> dict:
neighbor_list = neo4j_util.find_neighbors_with_relationships(
node_label=node_label,
node_properties=node_properties,
direction=direction,
rel_type=rel_type
)
# 提取所有唯一节点
node_dict = {}
for item in neighbor_list:
for key in ["source", "target"]:
n = item[key]
nid = n.get("id")
if nid and nid not in node_dict:
node_dict[nid] = n
# 如果没找到关系,但中心节点存在,也要包含它
if not neighbor_list:
center_nodes = neo4j_util.find_nodes_with_element_id(node_label, node_properties)
if center_nodes:
n = center_nodes[0]
node_dict[n["id"]] = n
nodes = list(node_dict.values())
relationships = neighbor_list # 结构已兼容
return build_g6_graph_data_from_results(nodes, relationships)
# ======================
# Robyn App
# ======================
app = Robyn(__file__)
@app.get("/getData")
def get_data():
try:
graph_data = build_g6_subgraph_by_props(
neo4j_client,
node_label="Disease",
node_properties={"name": "偏头痛"},
direction="both",
rel_type=None
)
return Response(
status_code=200,
description=jsonify(graph_data),
headers={"Content-Type": "text/plain; charset=utf-8"}
)
except Exception as e:
return jsonify({"error": str(e)}), 500
if __name__ == "__main__":
app.start(host="0.0.0.0", port=8088)