111 lines
3.7 KiB
Python
111 lines
3.7 KiB
Python
import os
|
|
import pickle
|
|
import networkx as nx
|
|
from typing import Optional
|
|
|
|
from enhanced_toon import EnhancedToon
|
|
|
|
class LocalGraph:
|
|
def __init__(self, root_path: str):
|
|
self.root_path = root_path
|
|
self.graph_dir = os.path.join(root_path, ".mcp_cache", "graph")
|
|
os.makedirs(self.graph_dir, exist_ok=True)
|
|
self.graph_path = os.path.join(self.graph_dir, "graph.pkl")
|
|
self.graph = nx.DiGraph()
|
|
|
|
# --- Creation ---
|
|
def add_node(self, node_id: str, **attrs):
|
|
self.graph.add_node(node_id, **attrs)
|
|
|
|
def add_edge(self, src: str, dst: str, edge_type: str, **attrs):
|
|
self.graph.add_edge(src, dst, type=edge_type, **attrs)
|
|
|
|
# --- Persistence ---
|
|
def save(self):
|
|
with open(self.graph_path, "wb") as f:
|
|
pickle.dump(self.graph, f, protocol=pickle.HIGHEST_PROTOCOL)
|
|
|
|
def load(self) -> bool:
|
|
if not os.path.exists(self.graph_path):
|
|
return False
|
|
with open(self.graph_path, "rb") as f:
|
|
self.graph = pickle.load(f)
|
|
return True
|
|
|
|
def clear(self):
|
|
self.graph = nx.DiGraph()
|
|
if os.path.exists(self.graph_path):
|
|
os.remove(self.graph_path)
|
|
|
|
# --- Query helpers ---
|
|
def find_nodes(self, **filters):
|
|
for node, data in self.graph.nodes(data=True):
|
|
if all(data.get(k) == v for k, v in filters.items()):
|
|
yield node, data
|
|
|
|
def neighbors(self, node_id, edge_type: Optional[str] = None):
|
|
for n in self.graph.successors(node_id):
|
|
if not edge_type or self.graph.edges[node_id, n]["type"] == edge_type:
|
|
yield n, self.graph.nodes[n]
|
|
|
|
def to_json(self):
|
|
out = {
|
|
"nodes": [
|
|
{"id": n, **data} for n, data in self.graph.nodes(data=True)
|
|
],
|
|
"edges": [
|
|
{"src": s, "dst": d, **data} for s, d, data in self.graph.edges(data=True)
|
|
],
|
|
}
|
|
path = os.path.join(self.graph_dir, "graph.json")
|
|
with open(path, "w") as f:
|
|
import json
|
|
json.dump(out, f, indent=2)
|
|
return path
|
|
|
|
def to_toon(self):
|
|
"""
|
|
Export graph to compact TOON (Token-Oriented Object Notation) format.
|
|
This reduces token count by 30-60% for LLM consumption.
|
|
"""
|
|
node_list = []
|
|
for node_id, data in self.graph.nodes(data=True):
|
|
# Build a flattened row
|
|
row = {
|
|
"id": node_id,
|
|
"type": data.get("type", ""),
|
|
"name": data.get("name", ""),
|
|
"lang": data.get("lang", ""),
|
|
"file": data.get("file", ""),
|
|
}
|
|
node_list.append(row)
|
|
|
|
edge_list = []
|
|
for src, dst, data in self.graph.edges(data=True):
|
|
edge_list.append({
|
|
"src": src,
|
|
"dst": dst,
|
|
"type": data.get("type", "")
|
|
})
|
|
|
|
lines = [
|
|
"# Graph Export",
|
|
f"nodes[{len(node_list)}]{{id,type,name,lang,file}}:"
|
|
]
|
|
for n in node_list:
|
|
row = [n["id"], n["type"], n["name"], n["lang"], n["file"]]
|
|
escaped = [EnhancedToon._escape_toon_field(str(f)) for f in row]
|
|
lines.append(" " + ",".join(escaped))
|
|
|
|
lines.append(f"edges[{len(edge_list)}]{{src,dst,type}}:")
|
|
for e in edge_list:
|
|
row = [e["src"], e["dst"], e["type"]]
|
|
escaped = [EnhancedToon._escape_toon_field(str(f)) for f in row]
|
|
lines.append(" " + ",".join(escaped))
|
|
|
|
toon_text = "\n".join(lines)
|
|
path = os.path.join(self.graph_dir, "graph.toon")
|
|
with open(path, "w", encoding="utf-8") as f:
|
|
f.write(toon_text)
|
|
return path
|