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 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:: 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