self-constructing_graph/graph/router/graph.py

from fastapi import APIRouter, Depends, HTTPException, status
from fastapi.security import OAuth2PasswordBearer, OAuth2PasswordRequestForm
from pydantic import BaseModel
from datetime import datetime, timedelta
from typing import Optional
from passlib.context import CryptContext
from db.neo4j import get_neo4j_db

# 路由
router = APIRouter()

def get_node_value(n):
    label = str(n.labels)
    if label == ":Disease":
        return 7000
    if label == ":Department":
        return 6000
    if label == ":Symptom":
        return 5500
    if label == ":Drug":
        return 5000
    return 4000
    
@router.get("/api/graph/{node_id}")
def get_graph_node_by_id(node_id: str):
    graph = get_neo4j_db()
    data = { "nodes":[], "edges":[]}    
    
    graph = get_neo4j_db()
  
    result = graph.run(f"match path=(n)-[*1]->(p) where id(n)={node_id} or id(p)={node_id} return nodes(path) as nodes,relationships(path) as relations limit 50")
    ids = []
    for r in result:
        nodes = r["nodes"]
        relations = r["relations"]
        
        for n in nodes:
            if n.identity in ids:
                continue
            else:
                ids.append(n.identity)
                data["nodes"].append({"id": n.identity, "name": n["name"], "value":get_node_value(n)})

        for r in relations:
            start_pos = ids.index(r.start_node.identity)
            end_pos = ids.index(r.end_node.identity)
            if start_pos >=0 and end_pos >=0:
                data["edges"].append({"id": r.identity, "name": type(r).__name__, "start": start_pos, "end":end_pos})
    return { "status": 200, "data": data}


@router.get("/api/node/{node_id}")
def get_node_by_id(node_id: str):
    graph = get_neo4j_db()
    data = { "nodes":[]}    
    
    graph = get_neo4j_db()
  
    result = graph.run(f"match (n) where id(n)={node_id} return n")
   
    for r in result:
        node = r["n"]
        itemData = {"id": node.identity,"type":str(node.labels), "name": node.get("name"), "value":5000, "props": []}
        for item in node.items():
            key, value = item
            if isinstance(value,list):
                value = ",".join(value)
            itemData["props"].append({ "name" : key, "value":value})
        
        data["nodes"].append(itemData)
            
    return { "status": 200, "data": data}


@router.get("/api/edge/{node_id}")
def get_edge_by_id(node_id: str):
    graph = get_neo4j_db()
    data = { "edges":[]}    
    
    graph = get_neo4j_db()
  
    result = graph.run(f"match (n)-[r]->(p) where id(r)={node_id} return r")
   
    for r in result:
        edges = r["r"]
        itemData = {"id": edges.identity,"type":type(edges).__name__, "name": edges.get("name"), "value":5000, "props": []}
        for item in edges.items():
            key, value = item
            if isinstance(value,list):
                value = ",".join(value)
            itemData["props"].append({ "name" : key, "value":value})
        
        data["edges"].append(itemData)
            
    return { "status": 200, "data": data}

@router.get("/api/graph-search/{name}")
def search_node_by_name(name: str):
    graph = get_neo4j_db()
    data = { "nodes":[], "edges":[]}    
    
    graph = get_neo4j_db()
  
    result = graph.run(f"match path=(n)-[*1]->(p) where n.name contains \"{name}\"  return nodes(path) as nodes,relationships(path) as relations limit 50")
    ids = []
    for r in result:
        nodes = r["nodes"]
        relations = r["relations"]
        
        for n in nodes:
            if n.identity in ids:
                continue
            else:
                ids.append(n.identity)
                data["nodes"].append({"id": n.identity, "name": n["name"], "value":get_node_value(n)})

        for r in relations:
            start_pos = ids.index(r.start_node.identity)
            end_pos = ids.index(r.end_node.identity)
            if start_pos >=0 and end_pos >=0:
                data["edges"].append({"id": r.identity, "name": type(r).__name__, "start": start_pos, "end":end_pos})
    return { "status": 200, "data": data}

graph_router = router