204 lines
6.3 KiB
Python
204 lines
6.3 KiB
Python
# -*- coding:utf-8 -*-
|
|
"""NetworkX-based graph engine with persistence."""
|
|
import os
|
|
import json
|
|
import networkx as nx
|
|
from typing import Dict, List, Any, Optional
|
|
|
|
DATA_DIR = "/data/ymq/graph-service/data"
|
|
_graphs: Dict[str, nx.DiGraph] = {}
|
|
|
|
|
|
def get_graph(graph_name: str = "default") -> nx.DiGraph:
|
|
"""Get or create a graph by name."""
|
|
if graph_name not in _graphs:
|
|
# Try to load from disk
|
|
filepath = os.path.join(DATA_DIR, f"{graph_name}.json")
|
|
if os.path.exists(filepath):
|
|
with open(filepath, 'r') as f:
|
|
data = json.load(f)
|
|
_graphs[graph_name] = nx.node_link_graph(data, directed=True)
|
|
else:
|
|
_graphs[graph_name] = nx.DiGraph()
|
|
return _graphs[graph_name]
|
|
|
|
|
|
def add_node(graph_name: str, node_id: str, **attrs) -> Dict:
|
|
"""Add a node to the graph."""
|
|
g = get_graph(graph_name)
|
|
g.add_node(node_id, **attrs)
|
|
return {"status": "ok", "node_id": node_id, "graph": graph_name}
|
|
|
|
|
|
def add_edge(graph_name: str, source: str, target: str, **attrs) -> Dict:
|
|
"""Add an edge to the graph."""
|
|
g = get_graph(graph_name)
|
|
g.add_edge(source, target, **attrs)
|
|
return {"status": "ok", "source": source, "target": target, "graph": graph_name}
|
|
|
|
|
|
def get_neighbors(graph_name: str, node_id: str, direction: str = "both", depth: int = 1) -> Dict:
|
|
"""Get neighbors of a node."""
|
|
g = get_graph(graph_name)
|
|
|
|
if node_id not in g:
|
|
return {"error": f"Node {node_id} not found"}
|
|
|
|
# BFS to get neighbors up to depth
|
|
visited = set()
|
|
queue = [(node_id, 0)]
|
|
neighbors = []
|
|
|
|
while queue:
|
|
current, d = queue.pop(0)
|
|
if d > depth:
|
|
break
|
|
|
|
if d > 0: # Don't include the starting node
|
|
neighbors.append({
|
|
"node_id": current,
|
|
"depth": d,
|
|
"attrs": dict(g.nodes[current])
|
|
})
|
|
|
|
if current not in visited:
|
|
visited.add(current)
|
|
|
|
if d < depth:
|
|
# Add successors
|
|
if direction in ["both", "out"]:
|
|
for neighbor in g.successors(current):
|
|
if neighbor not in visited:
|
|
queue.append((neighbor, d + 1))
|
|
|
|
# Add predecessors
|
|
if direction in ["both", "in"]:
|
|
for neighbor in g.predecessors(current):
|
|
if neighbor not in visited:
|
|
queue.append((neighbor, d + 1))
|
|
|
|
return {"node_id": node_id, "neighbors": neighbors, "count": len(neighbors)}
|
|
|
|
|
|
def find_path(graph_name: str, source: str, target: str, max_depth: int = 10) -> Dict:
|
|
"""Find shortest path between two nodes."""
|
|
g = get_graph(graph_name)
|
|
|
|
if source not in g:
|
|
return {"error": f"Source node {source} not found"}
|
|
if target not in g:
|
|
return {"error": f"Target node {target} not found"}
|
|
|
|
try:
|
|
path = nx.shortest_path(g, source, target)
|
|
if len(path) > max_depth + 1:
|
|
return {"error": f"Path too long: {len(path)-1} > {max_depth}"}
|
|
|
|
# Get path details
|
|
path_details = []
|
|
for i in range(len(path)):
|
|
node_id = path[i]
|
|
node_data = {"node_id": node_id, "attrs": dict(g.nodes[node_id])}
|
|
|
|
if i < len(path) - 1:
|
|
# Add edge info
|
|
next_node = path[i + 1]
|
|
edge_data = dict(g.edges[node_id, next_node])
|
|
node_data["edge_to_next"] = edge_data
|
|
|
|
path_details.append(node_data)
|
|
|
|
return {
|
|
"source": source,
|
|
"target": target,
|
|
"path": path,
|
|
"length": len(path) - 1,
|
|
"details": path_details
|
|
}
|
|
except nx.NetworkXNoPath:
|
|
return {"error": "No path found"}
|
|
|
|
|
|
def query_nodes(graph_name: str, filters: Optional[Dict] = None, limit: int = 100) -> Dict:
|
|
"""Query nodes with optional filters."""
|
|
g = get_graph(graph_name)
|
|
|
|
results = []
|
|
for node_id, attrs in g.nodes(data=True):
|
|
# Apply filters
|
|
match = True
|
|
if filters:
|
|
for key, value in filters.items():
|
|
if attrs.get(key) != value:
|
|
match = False
|
|
break
|
|
|
|
if match:
|
|
results.append({"node_id": node_id, "attrs": attrs})
|
|
if len(results) >= limit:
|
|
break
|
|
|
|
return {"nodes": results, "count": len(results)}
|
|
|
|
|
|
def get_stats(graph_name: str) -> Dict:
|
|
"""Get graph statistics."""
|
|
g = get_graph(graph_name)
|
|
|
|
stats = {
|
|
"graph": graph_name,
|
|
"nodes": g.number_of_nodes(),
|
|
"edges": g.number_of_edges(),
|
|
"density": nx.density(g),
|
|
"is_connected": nx.is_weakly_connected(g) if g.number_of_nodes() > 0 else True,
|
|
}
|
|
|
|
if g.number_of_nodes() > 0:
|
|
# Degree statistics
|
|
in_degrees = [d for _, d in g.in_degree()]
|
|
out_degrees = [d for _, d in g.out_degree()]
|
|
|
|
stats["avg_in_degree"] = sum(in_degrees) / len(in_degrees)
|
|
stats["avg_out_degree"] = sum(out_degrees) / len(out_degrees)
|
|
stats["max_in_degree"] = max(in_degrees)
|
|
stats["max_out_degree"] = max(out_degrees)
|
|
|
|
return stats
|
|
|
|
|
|
def save_graph(graph_name: str) -> Dict:
|
|
"""Save graph to disk."""
|
|
g = get_graph(graph_name)
|
|
filepath = os.path.join(DATA_DIR, f"{graph_name}.json")
|
|
|
|
data = nx.node_link_data(g)
|
|
with open(filepath, 'w') as f:
|
|
json.dump(data, f, ensure_ascii=False, indent=2)
|
|
|
|
return {"status": "ok", "filepath": filepath, "nodes": g.number_of_nodes(), "edges": g.number_of_edges()}
|
|
|
|
|
|
def load_graph(graph_name: str) -> Dict:
|
|
"""Load graph from disk."""
|
|
filepath = os.path.join(DATA_DIR, f"{graph_name}.json")
|
|
|
|
if not os.path.exists(filepath):
|
|
return {"error": f"Graph file not found: {filepath}"}
|
|
|
|
with open(filepath, 'r') as f:
|
|
data = json.load(f)
|
|
|
|
g = nx.node_link_graph(data, directed=True)
|
|
_graphs[graph_name] = g
|
|
|
|
return {"status": "ok", "nodes": g.number_of_nodes(), "edges": g.number_of_edges()}
|
|
|
|
|
|
def health_check():
|
|
"""Check service status."""
|
|
return {
|
|
"backend": "NetworkX",
|
|
"graphs_loaded": len(_graphs),
|
|
"data_dir": DATA_DIR
|
|
}
|