193 lines
6.9 KiB
Python
193 lines
6.9 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)
|
|
|
|
# --- 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
|
|
|
|
def _parse_node_id(self, node_id: str):
|
|
"""
|
|
Parse node_id into components and return a dict of canonical attributes.
|
|
|
|
Supported formats:
|
|
1) lang::node_type::file_path::name (full chunk/file tied)
|
|
2) lang::node_type::name (file-less canonical form - Option B)
|
|
3) file::<file_path> (file nodes)
|
|
otherwise -> fallback to Symbol
|
|
"""
|
|
attrs = {
|
|
"type": "Symbol",
|
|
"name": node_id,
|
|
"lang": "",
|
|
"file": ""
|
|
}
|
|
if not isinstance(node_id, str):
|
|
return attrs
|
|
parts = node_id.split("::")
|
|
if len(parts) == 4:
|
|
lang, node_type, file_path, name = parts
|
|
attrs["type"] = node_type
|
|
attrs["name"] = name
|
|
attrs["lang"] = lang
|
|
attrs["file"] = file_path
|
|
elif len(parts) == 3:
|
|
# Option B canonical file-less representation: lang::kind::entity_name
|
|
lang, node_type, name = parts
|
|
attrs["type"] = node_type
|
|
attrs["name"] = name
|
|
attrs["lang"] = lang
|
|
attrs["file"] = ""
|
|
elif len(parts) == 2 and parts[0] == "file":
|
|
# file::<file_path>
|
|
attrs["type"] = "File"
|
|
attrs["name"] = parts[1]
|
|
attrs["file"] = parts[1]
|
|
attrs["lang"] = ""
|
|
else:
|
|
# fallback: try to be helpful by guessing the name
|
|
if len(parts) >= 1:
|
|
attrs["name"] = parts[-1]
|
|
return attrs
|
|
|
|
def _ensure_node_exists(self, node_id: str):
|
|
"""
|
|
Ensure a node with node_id exists in the graph.
|
|
If it doesn't, create it using sensible attributes derived from the node_id format.
|
|
"""
|
|
if node_id in self.graph.nodes:
|
|
return
|
|
parsed = self._parse_node_id(node_id)
|
|
# Use add_node (keeps behaviour consistent)
|
|
self.add_node(node_id, **{
|
|
"type": parsed.get("type", "Symbol"),
|
|
"name": parsed.get("name", node_id),
|
|
"lang": parsed.get("lang", ""),
|
|
"file": parsed.get("file", "")
|
|
})
|
|
|
|
def add_edge(self, src: str, dst: str, edge_type: str, **attrs):
|
|
"""
|
|
Add an edge but first ensure both source and destination nodes exist and have basic attributes.
|
|
This prevents the creation of attribute-less nodes and makes semantic edges meaningful.
|
|
"""
|
|
try:
|
|
# Ensure source node exists (create minimal entry if missing)
|
|
if src not in self.graph.nodes:
|
|
self._ensure_node_exists(src)
|
|
|
|
# Ensure destination node exists (create minimal entry if missing)
|
|
if dst not in self.graph.nodes:
|
|
self._ensure_node_exists(dst)
|
|
|
|
# Finally add the edge with type and any extra attrs
|
|
self.graph.add_edge(src, dst, type=edge_type, **attrs)
|
|
except Exception:
|
|
# Keep behaviour non-fatal for indexing runs - log if you have a logger available
|
|
# fallback: still attempt to add the edge
|
|
try:
|
|
self.graph.add_edge(src, dst, type=edge_type, **attrs)
|
|
except Exception:
|
|
# swallow; graph should remain usable
|
|
pass
|