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): """Enhanced call edges that work cross-file by searching the entire graph.""" for cal in calls: if not cal or not isinstance(cal, str): continue # Try to find the called function in the graph target_found = False # Search for functions with this name across ALL files for node_id, data in self.graph.nodes(data=True): if (data.get('type') in ['function', 'method', 'async_function', 'constructor'] and data.get('name') == cal): # Found a matching function - create call edge self.add_edge(src, node_id, "calls") target_found = True # Don't break - a function might be called from multiple places # If not found, create a symbolic node for the call target if not target_found: # Create a canonical function node (file-less) for cross-file references symbolic_id = f"rust::function::{cal}" # Assuming Rust for now if symbolic_id not in self.graph.nodes: self.add_node( symbolic_id, type="function", name=cal, file="", # No specific file lang="rust", is_symbolic=True ) self.add_edge(src, symbolic_id, "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 def find_nodes_by_attributes(self, **filters): """Find nodes by attributes (not just exact matches).""" for node_id, data in self.graph.nodes(data=True): if all(data.get(k) == v for k, v in filters.items()): yield node_id, data def get_node_by_name_and_file(self, name: str, file_path: str, node_type: str = None, lang: str = None): """Find a node by name and file path with optional type/language filtering.""" candidates = [] for node_id, data in self.graph.nodes(data=True): if (data.get('name') == name and data.get('file') == file_path and (node_type is None or data.get('type') == node_type) and (lang is None or data.get('lang') == lang)): candidates.append((node_id, data)) return candidates def get_cross_file_relationships(self, node_id: str, current_file: str): """Get relationships that cross file boundaries.""" if node_id not in self.graph.nodes: return [], [] cross_file_out = [] cross_file_in = [] # Outgoing cross-file relationships for neighbor_id in self.graph.successors(node_id): if neighbor_id in self.graph.nodes: neighbor_file = self.graph.nodes[neighbor_id].get('file', '') if neighbor_file and neighbor_file != current_file: edge_data = self.graph.edges[node_id, neighbor_id] rel_type = edge_data.get('type', 'related') neighbor_name = self.graph.nodes[neighbor_id].get('name', '') neighbor_type = self.graph.nodes[neighbor_id].get('type', '') cross_file_out.append(f"{rel_type}:{neighbor_type}:{neighbor_name}@{neighbor_file}") # Incoming cross-file relationships for predecessor_id in self.graph.predecessors(node_id): if predecessor_id in self.graph.nodes: pred_file = self.graph.nodes[predecessor_id].get('file', '') if pred_file and pred_file != current_file: edge_data = self.graph.edges[predecessor_id, node_id] rel_type = edge_data.get('type', 'related') pred_name = self.graph.nodes[predecessor_id].get('name', '') pred_type = self.graph.nodes[predecessor_id].get('type', '') cross_file_in.append(f"{rel_type}:{pred_type}:{pred_name}@{pred_file}") return cross_file_out, cross_file_in def get_node_by_id(self, node_id: str) -> Optional[dict]: """Get node data by ID for graph queries.""" if node_id in self.graph.nodes: return dict(self.graph.nodes[node_id]) return None def get_relationships_for_query(self, node_id: str, max_relationships: int = 10) -> dict: """Get relationships in a format suitable for graph queries.""" if node_id not in self.graph.nodes: return {} result = { 'node_id': node_id, 'outgoing': [], 'incoming': [] } node_file = self.graph.nodes[node_id].get('file', '') # Outgoing relationships for neighbor_id in list(self.graph.successors(node_id))[:max_relationships]: if neighbor_id in self.graph.nodes: edge_data = self.graph.edges[node_id, neighbor_id] neighbor_data = self.graph.nodes[neighbor_id] rel_info = { 'type': edge_data.get('type', 'related'), 'target_id': neighbor_id, 'target_name': neighbor_data.get('name', ''), 'target_type': neighbor_data.get('type', ''), 'target_file': neighbor_data.get('file', ''), 'cross_file': neighbor_data.get('file', '') != node_file } result['outgoing'].append(rel_info) # Incoming relationships for predecessor_id in list(self.graph.predecessors(node_id))[:max_relationships]: if predecessor_id in self.graph.nodes: edge_data = self.graph.edges[predecessor_id, node_id] predecessor_data = self.graph.nodes[predecessor_id] rel_info = { 'type': edge_data.get('type', 'related'), 'source_id': predecessor_id, 'source_name': predecessor_data.get('name', ''), 'source_type': predecessor_data.get('type', ''), 'source_file': predecessor_data.get('file', ''), 'cross_file': predecessor_data.get('file', '') != node_file } result['incoming'].append(rel_info) return result def get_queryable_relationships(self, node_id: str) -> dict: """Get relationships in a format optimized for graph queries.""" if node_id not in self.graph.nodes: return {} node_data = self.graph.nodes[node_id] result = { 'node_id': node_id, 'node_type': node_data.get('type', ''), 'node_name': node_data.get('name', ''), 'file': node_data.get('file', ''), 'outgoing': [], 'incoming': [], 'cross_file': [], 'same_file': [] } node_file = node_data.get('file', '') # Process all relationships for neighbor_id in self.graph.successors(node_id): if neighbor_id in self.graph.nodes: edge_data = self.graph.edges[node_id, neighbor_id] neighbor_data = self.graph.nodes[neighbor_id] rel_info = { 'type': edge_data.get('type', 'related'), 'target_id': neighbor_id, 'target_name': neighbor_data.get('name', ''), 'target_type': neighbor_data.get('type', ''), 'target_file': neighbor_data.get('file', '') } result['outgoing'].append(rel_info) # Classify by file relationship if neighbor_data.get('file') and neighbor_data.get('file') != node_file: result['cross_file'].append(rel_info) else: result['same_file'].append(rel_info) # Reverse relationships for predecessor_id in self.graph.predecessors(node_id): if predecessor_id in self.graph.nodes: edge_data = self.graph.edges[predecessor_id, node_id] predecessor_data = self.graph.nodes[predecessor_id] rel_info = { 'type': edge_data.get('type', 'related'), 'source_id': predecessor_id, 'source_name': predecessor_data.get('name', ''), 'source_type': predecessor_data.get('type', ''), 'source_file': predecessor_data.get('file', '') } result['incoming'].append(rel_info) return result def classify_relationships(self, node_id): node_data = self.graph.nodes[node_id] node_file = node_data.get('file', '') result = { 'outgoing': [], 'incoming': [], 'same_file': [], 'cross_file': [] } for neighbor_id in self.graph.successors(node_id): if neighbor_id in self.graph.nodes: edge_data = self.graph.edges[node_id, neighbor_id] neighbor_data = self.graph.nodes[neighbor_id] rel_info = { 'type': edge_data.get('type', 'related'), 'target_id': neighbor_id, 'target_name': neighbor_data.get('name', ''), 'target_type': neighbor_data.get('type', ''), 'target_file': neighbor_data.get('file', '') } result['outgoing'].append(rel_info) # Classify by file relationship if neighbor_data.get('file') and neighbor_data.get('file') != node_file: result['cross_file'].append(rel_info) else: result['same_file'].append(rel_info) # Reverse relationships for predecessor_id in self.graph.predecessors(node_id): if predecessor_id in self.graph.nodes: edge_data = self.graph.edges[predecessor_id, node_id] predecessor_data = self.graph.nodes[predecessor_id] rel_info = { 'type': edge_data.get('type', 'related'), 'source_id': predecessor_id, 'source_name': predecessor_data.get('name', ''), 'source_type': predecessor_data.get('type', ''), 'source_file': predecessor_data.get('file', '') } result['incoming'].append(rel_info) return result