mirror of
https://github.com/Graphify-Labs/graphify.git
synced 2026-09-14 19:34:09 +08:00
Safi
489 lines
20 KiB
Python
489 lines
20 KiB
Python
# MCP stdio server - exposes graph query tools to Claude and other agents
|
|
from __future__ import annotations
|
|
import json
|
|
import sys
|
|
from pathlib import Path
|
|
import networkx as nx
|
|
from networkx.readwrite import json_graph
|
|
from graphify.security import sanitize_label
|
|
|
|
|
|
def _load_graph(graph_path: str) -> nx.Graph:
|
|
try:
|
|
resolved = Path(graph_path).resolve()
|
|
if resolved.suffix != ".json":
|
|
raise ValueError(f"Graph path must be a .json file, got: {graph_path!r}")
|
|
if not resolved.exists():
|
|
raise FileNotFoundError(f"Graph file not found: {resolved}")
|
|
safe = resolved
|
|
data = json.loads(safe.read_text(encoding="utf-8"))
|
|
try:
|
|
return json_graph.node_link_graph(data, edges="links")
|
|
except TypeError:
|
|
return json_graph.node_link_graph(data)
|
|
except (ValueError, FileNotFoundError) as exc:
|
|
print(f"error: {exc}", file=sys.stderr)
|
|
sys.exit(1)
|
|
except json.JSONDecodeError as exc:
|
|
print(f"error: graph.json is corrupted ({exc}). Re-run /graphify to rebuild.", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
|
|
def _communities_from_graph(G: nx.Graph) -> dict[int, list[str]]:
|
|
"""Reconstruct community dict from community property stored on nodes."""
|
|
communities: dict[int, list[str]] = {}
|
|
for node_id, data in G.nodes(data=True):
|
|
cid = data.get("community")
|
|
if cid is not None:
|
|
communities.setdefault(int(cid), []).append(node_id)
|
|
return communities
|
|
|
|
|
|
def _strip_diacritics(text: str) -> str:
|
|
import unicodedata
|
|
nfkd = unicodedata.normalize("NFKD", text)
|
|
return "".join(c for c in nfkd if not unicodedata.combining(c))
|
|
|
|
|
|
_EXACT_MATCH_BONUS = 100.0
|
|
|
|
|
|
def _score_nodes(G: nx.Graph, terms: list[str]) -> list[tuple[float, str]]:
|
|
scored = []
|
|
norm_terms = [_strip_diacritics(t).lower() for t in terms]
|
|
for nid, data in G.nodes(data=True):
|
|
norm_label = data.get("norm_label") or _strip_diacritics(data.get("label") or "").lower()
|
|
source = (data.get("source_file") or "").lower()
|
|
score = sum(1 for t in norm_terms if t in norm_label) + sum(0.5 for t in norm_terms if t in source)
|
|
# Exact match: single term equals the full label (strip trailing () for functions)
|
|
if any(t == norm_label or t == norm_label.rstrip("()") for t in norm_terms):
|
|
score += _EXACT_MATCH_BONUS
|
|
if score > 0:
|
|
scored.append((score, nid))
|
|
return sorted(scored, reverse=True)
|
|
|
|
|
|
_CONTEXT_HINTS: tuple[tuple[str, tuple[str, ...]], ...] = (
|
|
("call", ("call", "calls", "called", "invoke", "invokes", "invoked")),
|
|
("import", ("import", "imports", "imported", "module", "modules")),
|
|
("field", ("field", "fields", "member", "members", "property", "properties")),
|
|
("parameter_type", ("parameter", "parameters", "param", "params", "argument", "arguments")),
|
|
("return_type", ("return", "returns", "returned")),
|
|
("generic_arg", ("generic", "generics", "template", "templates")),
|
|
)
|
|
|
|
|
|
def _normalize_context_filters(filters: list[str] | None) -> list[str]:
|
|
if not filters:
|
|
return []
|
|
normalized: list[str] = []
|
|
seen: set[str] = set()
|
|
for value in filters:
|
|
key = _strip_diacritics(str(value)).strip().lower()
|
|
if key and key not in seen:
|
|
seen.add(key)
|
|
normalized.append(key)
|
|
return normalized
|
|
|
|
|
|
def _infer_context_filters(question: str) -> list[str]:
|
|
lowered = {
|
|
_strip_diacritics(token).lower()
|
|
for token in question.replace("?", " ").replace(",", " ").split()
|
|
}
|
|
inferred: list[str] = []
|
|
for context, hints in _CONTEXT_HINTS:
|
|
if any(hint in lowered for hint in hints):
|
|
inferred.append(context)
|
|
return inferred
|
|
|
|
|
|
def _resolve_context_filters(question: str, explicit_filters: list[str] | None = None) -> tuple[list[str], str | None]:
|
|
normalized = _normalize_context_filters(explicit_filters)
|
|
if normalized:
|
|
return normalized, "explicit"
|
|
inferred = _infer_context_filters(question)
|
|
if inferred:
|
|
return inferred, "heuristic"
|
|
return [], None
|
|
|
|
|
|
def _filter_graph_by_context(G: nx.Graph, context_filters: list[str] | None) -> nx.Graph:
|
|
filters = set(_normalize_context_filters(context_filters))
|
|
if not filters:
|
|
return G
|
|
H = G.__class__()
|
|
H.add_nodes_from(G.nodes(data=True))
|
|
if isinstance(G, (nx.MultiGraph, nx.MultiDiGraph)):
|
|
for u, v, key, data in G.edges(keys=True, data=True):
|
|
if data.get("context") in filters:
|
|
H.add_edge(u, v, key=key, **data)
|
|
else:
|
|
for u, v, data in G.edges(data=True):
|
|
if data.get("context") in filters:
|
|
H.add_edge(u, v, **data)
|
|
return H
|
|
|
|
|
|
def _bfs(G: nx.Graph, start_nodes: list[str], depth: int) -> tuple[set[str], list[tuple]]:
|
|
visited: set[str] = set(start_nodes)
|
|
frontier = set(start_nodes)
|
|
edges_seen: list[tuple] = []
|
|
for _ in range(depth):
|
|
next_frontier: set[str] = set()
|
|
for n in frontier:
|
|
for neighbor in G.neighbors(n):
|
|
if neighbor not in visited:
|
|
next_frontier.add(neighbor)
|
|
edges_seen.append((n, neighbor))
|
|
visited.update(next_frontier)
|
|
frontier = next_frontier
|
|
return visited, edges_seen
|
|
|
|
|
|
def _dfs(G: nx.Graph, start_nodes: list[str], depth: int) -> tuple[set[str], list[tuple]]:
|
|
visited: set[str] = set()
|
|
edges_seen: list[tuple] = []
|
|
stack = [(n, 0) for n in reversed(start_nodes)]
|
|
while stack:
|
|
node, d = stack.pop()
|
|
if node in visited or d > depth:
|
|
continue
|
|
visited.add(node)
|
|
for neighbor in G.neighbors(node):
|
|
if neighbor not in visited:
|
|
stack.append((neighbor, d + 1))
|
|
edges_seen.append((node, neighbor))
|
|
return visited, edges_seen
|
|
|
|
|
|
def _subgraph_to_text(G: nx.Graph, nodes: set[str], edges: list[tuple], token_budget: int = 2000, *, seeds: list[str] | None = None) -> str:
|
|
"""Render subgraph as text, cutting at token_budget (approx 3 chars/token).
|
|
|
|
seeds: exact-match nodes rendered first before the degree-sorted expansion,
|
|
so the queried symbol always appears at the top of the output.
|
|
"""
|
|
char_budget = token_budget * 3
|
|
lines = []
|
|
seed_set = set(seeds or [])
|
|
ordered = [n for n in (seeds or []) if n in nodes] + \
|
|
sorted(nodes - seed_set, key=lambda n: G.degree(n), reverse=True)
|
|
for nid in ordered:
|
|
d = G.nodes[nid]
|
|
line = f"NODE {sanitize_label(d.get('label', nid))} [src={d.get('source_file', '')} loc={d.get('source_location', '')} community={d.get('community', '')}]"
|
|
lines.append(line)
|
|
for u, v in edges:
|
|
if u in nodes and v in nodes:
|
|
raw = G[u][v]
|
|
d = next(iter(raw.values()), {}) if isinstance(G, (nx.MultiGraph, nx.MultiDiGraph)) else raw
|
|
context = d.get("context")
|
|
context_suffix = f" context={context}" if context else ""
|
|
line = (
|
|
f"EDGE {sanitize_label(G.nodes[u].get('label', u))} "
|
|
f"--{d.get('relation', '')} [{d.get('confidence', '')}{context_suffix}]--> "
|
|
f"{sanitize_label(G.nodes[v].get('label', v))}"
|
|
)
|
|
lines.append(line)
|
|
output = "\n".join(lines)
|
|
if len(output) > char_budget:
|
|
output = output[:char_budget] + f"\n... (truncated to ~{token_budget} token budget)"
|
|
return output
|
|
|
|
|
|
def _query_graph_text(
|
|
G: nx.Graph,
|
|
question: str,
|
|
*,
|
|
mode: str = "bfs",
|
|
depth: int = 3,
|
|
token_budget: int = 2000,
|
|
context_filters: list[str] | None = None,
|
|
) -> str:
|
|
terms = [t.lower() for t in question.split() if len(t) > 2]
|
|
scored = _score_nodes(G, terms)
|
|
start_nodes = [nid for _, nid in scored[:3]]
|
|
if not start_nodes:
|
|
return "No matching nodes found."
|
|
resolved_filters, filter_source = _resolve_context_filters(question, context_filters)
|
|
traversal_graph = _filter_graph_by_context(G, resolved_filters)
|
|
nodes, edges = _dfs(traversal_graph, start_nodes, depth) if mode == "dfs" else _bfs(traversal_graph, start_nodes, depth)
|
|
header_parts = [
|
|
f"Traversal: {mode.upper()} depth={depth}",
|
|
f"Start: {[G.nodes[n].get('label', n) for n in start_nodes]}",
|
|
]
|
|
if resolved_filters:
|
|
header_parts.append(f"Context: {', '.join(resolved_filters)} ({filter_source})")
|
|
header_parts.append(f"{len(nodes)} nodes found")
|
|
header = " | ".join(header_parts) + "\n\n"
|
|
return header + _subgraph_to_text(traversal_graph, nodes, edges, token_budget)
|
|
|
|
|
|
def _find_node(G: nx.Graph, label: str) -> list[str]:
|
|
"""Return node IDs whose label or ID matches the search term (diacritic-insensitive)."""
|
|
term = _strip_diacritics(label).lower()
|
|
return [nid for nid, d in G.nodes(data=True)
|
|
if term in (d.get("norm_label") or _strip_diacritics(d.get("label") or "").lower())
|
|
or term == nid.lower()]
|
|
|
|
|
|
def _filter_blank_stdin() -> None:
|
|
"""Filter blank lines from stdin before MCP reads it.
|
|
|
|
Some MCP clients (Claude Desktop, etc.) send blank lines between JSON
|
|
messages. The MCP stdio transport tries to parse every line as a
|
|
JSONRPCMessage, so a bare newline triggers a Pydantic ValidationError.
|
|
This installs an OS-level pipe that relays stdin while dropping blanks.
|
|
"""
|
|
import os
|
|
import threading
|
|
|
|
r_fd, w_fd = os.pipe()
|
|
saved_fd = os.dup(sys.stdin.fileno())
|
|
|
|
def _relay() -> None:
|
|
try:
|
|
with open(saved_fd, "rb") as src, open(w_fd, "wb") as dst:
|
|
for line in src:
|
|
if line.strip():
|
|
dst.write(line)
|
|
dst.flush()
|
|
except Exception:
|
|
pass
|
|
|
|
threading.Thread(target=_relay, daemon=True).start()
|
|
os.dup2(r_fd, sys.stdin.fileno())
|
|
os.close(r_fd)
|
|
sys.stdin = open(0, "r", closefd=False)
|
|
|
|
|
|
def serve(graph_path: str = "graphify-out/graph.json") -> None:
|
|
"""Start the MCP server. Requires pip install mcp."""
|
|
try:
|
|
from mcp.server import Server
|
|
from mcp.server.stdio import stdio_server
|
|
from mcp import types
|
|
except ImportError as e:
|
|
raise ImportError("mcp not installed. Run: pip install mcp") from e
|
|
|
|
G = _load_graph(graph_path)
|
|
communities = _communities_from_graph(G)
|
|
|
|
server = Server("graphify")
|
|
|
|
@server.list_tools()
|
|
async def list_tools() -> list[types.Tool]:
|
|
return [
|
|
types.Tool(
|
|
name="query_graph",
|
|
description="Search the knowledge graph using BFS or DFS. Returns relevant nodes and edges as text context.",
|
|
inputSchema={
|
|
"type": "object",
|
|
"properties": {
|
|
"question": {"type": "string", "description": "Natural language question or keyword search"},
|
|
"mode": {"type": "string", "enum": ["bfs", "dfs"], "default": "bfs",
|
|
"description": "bfs=broad context, dfs=trace a specific path"},
|
|
"depth": {"type": "integer", "default": 3, "description": "Traversal depth (1-6)"},
|
|
"token_budget": {"type": "integer", "default": 2000, "description": "Max output tokens"},
|
|
"context_filter": {
|
|
"type": "array",
|
|
"items": {"type": "string"},
|
|
"description": "Optional explicit edge-context filter, e.g. ['call', 'field']",
|
|
},
|
|
},
|
|
"required": ["question"],
|
|
},
|
|
),
|
|
types.Tool(
|
|
name="get_node",
|
|
description="Get full details for a specific node by label or ID.",
|
|
inputSchema={
|
|
"type": "object",
|
|
"properties": {"label": {"type": "string", "description": "Node label or ID to look up"}},
|
|
"required": ["label"],
|
|
},
|
|
),
|
|
types.Tool(
|
|
name="get_neighbors",
|
|
description="Get all direct neighbors of a node with edge details.",
|
|
inputSchema={
|
|
"type": "object",
|
|
"properties": {
|
|
"label": {"type": "string"},
|
|
"relation_filter": {"type": "string", "description": "Optional: filter by relation type"},
|
|
},
|
|
"required": ["label"],
|
|
},
|
|
),
|
|
types.Tool(
|
|
name="get_community",
|
|
description="Get all nodes in a community by community ID.",
|
|
inputSchema={
|
|
"type": "object",
|
|
"properties": {"community_id": {"type": "integer", "description": "Community ID (0-indexed by size)"}},
|
|
"required": ["community_id"],
|
|
},
|
|
),
|
|
types.Tool(
|
|
name="god_nodes",
|
|
description="Return the most connected nodes - the core abstractions of the knowledge graph.",
|
|
inputSchema={"type": "object", "properties": {"top_n": {"type": "integer", "default": 10}}},
|
|
),
|
|
types.Tool(
|
|
name="graph_stats",
|
|
description="Return summary statistics: node count, edge count, communities, confidence breakdown.",
|
|
inputSchema={"type": "object", "properties": {}},
|
|
),
|
|
types.Tool(
|
|
name="shortest_path",
|
|
description="Find the shortest path between two concepts in the knowledge graph.",
|
|
inputSchema={
|
|
"type": "object",
|
|
"properties": {
|
|
"source": {"type": "string", "description": "Source concept label or keyword"},
|
|
"target": {"type": "string", "description": "Target concept label or keyword"},
|
|
"max_hops": {"type": "integer", "default": 8, "description": "Maximum hops to consider"},
|
|
},
|
|
"required": ["source", "target"],
|
|
},
|
|
),
|
|
]
|
|
|
|
def _tool_query_graph(arguments: dict) -> str:
|
|
question = arguments["question"]
|
|
mode = arguments.get("mode", "bfs")
|
|
depth = min(int(arguments.get("depth", 3)), 6)
|
|
budget = int(arguments.get("token_budget", 2000))
|
|
context_filter = arguments.get("context_filter")
|
|
return _query_graph_text(
|
|
G,
|
|
question,
|
|
mode=mode,
|
|
depth=depth,
|
|
token_budget=budget,
|
|
context_filters=context_filter,
|
|
)
|
|
|
|
def _tool_get_node(arguments: dict) -> str:
|
|
label = arguments["label"].lower()
|
|
matches = [(nid, d) for nid, d in G.nodes(data=True)
|
|
if label in (d.get("label") or "").lower() or label == nid.lower()]
|
|
if not matches:
|
|
return f"No node matching '{label}' found."
|
|
nid, d = matches[0]
|
|
return "\n".join([
|
|
f"Node: {d.get('label', nid)}",
|
|
f" ID: {nid}",
|
|
f" Source: {d.get('source_file', '')} {d.get('source_location', '')}",
|
|
f" Type: {d.get('file_type', '')}",
|
|
f" Community: {d.get('community', '')}",
|
|
f" Degree: {G.degree(nid)}",
|
|
])
|
|
|
|
def _tool_get_neighbors(arguments: dict) -> str:
|
|
label = arguments["label"].lower()
|
|
rel_filter = arguments.get("relation_filter", "").lower()
|
|
matches = _find_node(G, label)
|
|
if not matches:
|
|
return f"No node matching '{label}' found."
|
|
nid = matches[0]
|
|
lines = [f"Neighbors of {G.nodes[nid].get('label', nid)}:"]
|
|
for neighbor in G.neighbors(nid):
|
|
d = G.edges[nid, neighbor]
|
|
rel = d.get("relation", "")
|
|
if rel_filter and rel_filter not in rel.lower():
|
|
continue
|
|
lines.append(f" --> {G.nodes[neighbor].get('label', neighbor)} [{rel}] [{d.get('confidence', '')}]")
|
|
return "\n".join(lines)
|
|
|
|
def _tool_get_community(arguments: dict) -> str:
|
|
cid = int(arguments["community_id"])
|
|
nodes = communities.get(cid, [])
|
|
if not nodes:
|
|
return f"Community {cid} not found."
|
|
lines = [f"Community {cid} ({len(nodes)} nodes):"]
|
|
for n in nodes:
|
|
d = G.nodes[n]
|
|
lines.append(f" {d.get('label', n)} [{d.get('source_file', '')}]")
|
|
return "\n".join(lines)
|
|
|
|
def _tool_god_nodes(arguments: dict) -> str:
|
|
from .analyze import god_nodes as _god_nodes
|
|
nodes = _god_nodes(G, top_n=int(arguments.get("top_n", 10)))
|
|
lines = ["God nodes (most connected):"]
|
|
lines += [f" {i}. {n['label']} - {n['degree']} edges" for i, n in enumerate(nodes, 1)]
|
|
return "\n".join(lines)
|
|
|
|
def _tool_graph_stats(_: dict) -> str:
|
|
confs = [d.get("confidence", "EXTRACTED") for _, _, d in G.edges(data=True)]
|
|
total = len(confs) or 1
|
|
return (
|
|
f"Nodes: {G.number_of_nodes()}\n"
|
|
f"Edges: {G.number_of_edges()}\n"
|
|
f"Communities: {len(communities)}\n"
|
|
f"EXTRACTED: {round(confs.count('EXTRACTED')/total*100)}%\n"
|
|
f"INFERRED: {round(confs.count('INFERRED')/total*100)}%\n"
|
|
f"AMBIGUOUS: {round(confs.count('AMBIGUOUS')/total*100)}%\n"
|
|
)
|
|
|
|
def _tool_shortest_path(arguments: dict) -> str:
|
|
src_scored = _score_nodes(G, [t.lower() for t in arguments["source"].split()])
|
|
tgt_scored = _score_nodes(G, [t.lower() for t in arguments["target"].split()])
|
|
if not src_scored:
|
|
return f"No node matching source '{arguments['source']}' found."
|
|
if not tgt_scored:
|
|
return f"No node matching target '{arguments['target']}' found."
|
|
src_nid, tgt_nid = src_scored[0][1], tgt_scored[0][1]
|
|
max_hops = int(arguments.get("max_hops", 8))
|
|
try:
|
|
path_nodes = nx.shortest_path(G, src_nid, tgt_nid)
|
|
except (nx.NetworkXNoPath, nx.NodeNotFound):
|
|
return f"No path found between '{G.nodes[src_nid].get('label', src_nid)}' and '{G.nodes[tgt_nid].get('label', tgt_nid)}'."
|
|
hops = len(path_nodes) - 1
|
|
if hops > max_hops:
|
|
return f"Path exceeds max_hops={max_hops} ({hops} hops found)."
|
|
segments = []
|
|
for i in range(len(path_nodes) - 1):
|
|
u, v = path_nodes[i], path_nodes[i + 1]
|
|
edata = G.edges[u, v]
|
|
rel = edata.get("relation", "")
|
|
conf = edata.get("confidence", "")
|
|
conf_str = f" [{conf}]" if conf else ""
|
|
if i == 0:
|
|
segments.append(G.nodes[u].get("label", u))
|
|
segments.append(f"--{rel}{conf_str}--> {G.nodes[v].get('label', v)}")
|
|
return f"Shortest path ({hops} hops):\n " + " ".join(segments)
|
|
|
|
_handlers = {
|
|
"query_graph": _tool_query_graph,
|
|
"get_node": _tool_get_node,
|
|
"get_neighbors": _tool_get_neighbors,
|
|
"get_community": _tool_get_community,
|
|
"god_nodes": _tool_god_nodes,
|
|
"graph_stats": _tool_graph_stats,
|
|
"shortest_path": _tool_shortest_path,
|
|
}
|
|
|
|
@server.call_tool()
|
|
async def call_tool(name: str, arguments: dict) -> list[types.TextContent]:
|
|
handler = _handlers.get(name)
|
|
if not handler:
|
|
return [types.TextContent(type="text", text=f"Unknown tool: {name}")]
|
|
try:
|
|
return [types.TextContent(type="text", text=handler(arguments))]
|
|
except Exception as exc:
|
|
return [types.TextContent(type="text", text=f"Error executing {name}: {exc}")]
|
|
|
|
import asyncio
|
|
|
|
async def main() -> None:
|
|
async with stdio_server() as streams:
|
|
await server.run(streams[0], streams[1], server.create_initialization_options())
|
|
|
|
_filter_blank_stdin()
|
|
asyncio.run(main())
|
|
|
|
|
|
if __name__ == "__main__":
|
|
graph_path = sys.argv[1] if len(sys.argv) > 1 else "graphify-out/graph.json"
|
|
serve(graph_path)
|