287 lines
10 KiB
Python
287 lines
10 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 | dict, **attrs):
|
|
"""
|
|
Enhanced add_node that can return the generated node_id
|
|
"""
|
|
if isinstance(node_id, dict):
|
|
lang = node_id.get("lang", "unknown")
|
|
kind = node_id.get("type", "Symbol")
|
|
name = node_id.get("name", "unknown")
|
|
fpath = node_id.get("file", "")
|
|
node_id = f"{lang}::{kind}::{fpath}::{name}"
|
|
|
|
# Merge any additional attributes
|
|
final_attrs = {}
|
|
if isinstance(node_id, dict):
|
|
final_attrs.update(node_id)
|
|
final_attrs.update(attrs)
|
|
|
|
self.graph.add_node(node_id, **final_attrs)
|
|
return node_id # Return the ID for reference
|
|
|
|
def _add_import_edges(self, src: str, imports: list[str]):
|
|
for imp in imports:
|
|
# Resolve import to a file node if present
|
|
target = f"file::{imp}"
|
|
self.add_edge(src, target, "imports")
|
|
|
|
def _add_call_edges(self, src: str, calls: list[str], file_path: str):
|
|
for cal in calls:
|
|
target = f"rust::function::{file_path}::{cal}"
|
|
self.add_edge(src, target, "calls")
|
|
|
|
|
|
|
|
# --- 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
|
|
|
|
def get_node_relationships(self, node_id: str) -> dict:
|
|
"""
|
|
Get comprehensive relationships for a node in a structured format.
|
|
Returns: {
|
|
'outgoing': [(neighbor_id, edge_type, neighbor_data)],
|
|
'incoming': [(predecessor_id, edge_type, predecessor_data)],
|
|
'cross_file': list of cross-file relationships,
|
|
'same_file': list of same-file relationships
|
|
}
|
|
"""
|
|
if node_id not in self.graph:
|
|
return {}
|
|
|
|
result = {
|
|
'outgoing': [],
|
|
'incoming': [],
|
|
'cross_file': [],
|
|
'same_file': []
|
|
}
|
|
|
|
node_file = self.graph.nodes[node_id].get('file', '')
|
|
|
|
# Outgoing relationships
|
|
for neighbor_id in self.graph.successors(node_id):
|
|
edge_data = self.graph.edges[node_id, neighbor_id]
|
|
neighbor_data = self.graph.nodes[neighbor_id]
|
|
edge_type = edge_data.get('type', 'related')
|
|
|
|
result['outgoing'].append((neighbor_id, edge_type, neighbor_data))
|
|
|
|
# Cross-file classification
|
|
neighbor_file = neighbor_data.get('file', '')
|
|
rel_info = f"{edge_type}:{neighbor_data.get('type', '')}:{neighbor_data.get('name', '')}"
|
|
|
|
if neighbor_file and neighbor_file != node_file:
|
|
result['cross_file'].append(f"{rel_info}@{neighbor_file}")
|
|
else:
|
|
result['same_file'].append(rel_info)
|
|
|
|
# Incoming relationships
|
|
for predecessor_id in self.graph.predecessors(node_id):
|
|
edge_data = self.graph.edges[predecessor_id, node_id]
|
|
predecessor_data = self.graph.nodes[predecessor_id]
|
|
edge_type = edge_data.get('type', 'related')
|
|
|
|
result['incoming'].append((predecessor_id, edge_type, predecessor_data))
|
|
|
|
return result
|
|
|
|
def get_file_entities(self, file_path: str) -> list:
|
|
"""Get all entities in a file in compact format."""
|
|
file_node_id = f"file::{file_path}"
|
|
if file_node_id not in self.graph:
|
|
return []
|
|
|
|
entities = []
|
|
for neighbor_id, neighbor_data in self.neighbors(file_node_id):
|
|
entity_type = neighbor_data.get('type', '')
|
|
entity_name = neighbor_data.get('name', '')
|
|
if entity_type and entity_name:
|
|
entities.append(f"{entity_type}:{entity_name}")
|
|
|
|
return entities
|