From 76bab8abc244dde6a4e76f2c443d3cc104aac13a Mon Sep 17 00:00:00 2001 From: Stephen Adamson Date: Sat, 15 Nov 2025 23:48:16 +0000 Subject: [PATCH] update graph --- graph/graph.py | 81 ++++++++++++++++- mcp_codebase.py | 224 ++++++++++++++++++++---------------------------- 2 files changed, 172 insertions(+), 133 deletions(-) diff --git a/graph/graph.py b/graph/graph.py index 0e9dd3d..f47d342 100644 --- a/graph/graph.py +++ b/graph/graph.py @@ -41,9 +41,37 @@ class LocalGraph: 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: - target = f"rust::function::{file_path}::{cal}" - self.add_edge(src, target, "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") @@ -284,3 +312,52 @@ class LocalGraph: 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 diff --git a/mcp_codebase.py b/mcp_codebase.py index e1ff29a..305d21c 100644 --- a/mcp_codebase.py +++ b/mcp_codebase.py @@ -2865,6 +2865,7 @@ def build_indexes(): trait_registry = {} class_registry = {} package_registry = {} + function_registry = {} logger.info("=== Phase 1: Creating nodes ===") @@ -2951,6 +2952,14 @@ def build_indexes(): class_node_id = f"{lang}::class::{fpath}::{class_name}" graph.add_edge(class_node_id, node_id, "contains") + + # Register in function_registry for cross-file calls + function_registry[name] = { + "node_id": node_id, + "lang": lang, + "file": fpath, + "type": node_type + } else: graph.add_node( node_id, @@ -3317,8 +3326,18 @@ def build_indexes(): calls = calls_raw if isinstance(calls_raw, list) else [] if calls: + # Use the enhanced _add_call_edges that handles cross-file graph._add_call_edges(node_id, calls, fpath) + # ALSO create direct edges to functions we can find in registries + for cal in calls: + # Search function_registry if we had one, but for now search created_nodes + for candidate_id, candidate_data in created_nodes.items(): + if (candidate_data.get('type') in ['function', 'method'] and + candidate_data.get('name') == cal and + candidate_data.get('file') != fpath): # Cross-file only + graph.add_edge(node_id, candidate_id, "calls") + # ===== CLASS/STRUCT INHERITANCE ===== if node_type in ["class", "struct"]: extends = m.get("extends") @@ -3668,7 +3687,6 @@ def _extract_param_type(param: str) -> str: return _clean_type_string(param_type) - def _extract_field_type(field: str) -> str: """Extract type from field string for all languages.""" if not field or not isinstance(field, str): @@ -3698,7 +3716,6 @@ def _extract_field_type(field: str) -> str: return _clean_type_string(field_type) - def build_cross_language_relationships(graph, created_nodes, component_registry, interface_registry, class_registry): """Build relationships across different programming languages.""" @@ -3744,7 +3761,6 @@ def build_cross_language_relationships(graph, created_nodes, component_registry, return connections - def build_type_relationships(graph, created_nodes): """Build enhanced type relationships within language boundaries.""" connections = 0 @@ -3823,7 +3839,6 @@ def build_type_relationships(graph, created_nodes): return connections - def build_bevy_relationships(graph, created_nodes): """Build Bevy-specific relationships (ECS patterns).""" connections = 0 @@ -3907,102 +3922,6 @@ def build_bevy_relationships(graph, created_nodes): return connections -def build_bevy_relationships(graph, created_nodes): - """Build Bevy-specific relationships (ECS patterns).""" - connections = 0 - - # Track Bevy systems, components, and resources - systems = [] - components = [] - resources = [] - bundles = [] - - for node_id, node_data in created_nodes.items(): - m = node_data.get("metadata", {}) - if not m: - continue - - lang = node_data.get("lang") - if lang != "rust": - continue - - node_type = node_data.get("type") - name = node_data.get("name") - - # Detect Bevy systems (functions with Query parameters) - if node_type in ["function", "method"]: - params = m.get("parameters", []) - has_query = any("Query" in str(p) for p in params) - has_commands = any("Commands" in str(p) for p in params) - has_res = any("Res" in str(p) or "ResMut" in str(p) for p in params) - - if has_query or has_commands or has_res: - systems.append(node_id) - - # Detect Bevy components (structs with Component derive) - if node_type == "struct": - derives = m.get("derives", []) - if "Component" in str(derives): - components.append(node_id) - - # Detect Bevy resources (structs with Resource derive) - if node_type == "struct": - derives = m.get("derives", []) - if "Resource" in str(derives): - resources.append(node_id) - - # Detect Bevy bundles (structs with Bundle derive) - if node_type == "struct": - derives = m.get("derives", []) - if "Bundle" in str(derives): - bundles.append(node_id) - - # Connect systems to the components/resources they use - for system_id in systems: - system_data = created_nodes.get(system_id, {}) - m = system_data.get("metadata", {}) - params = m.get("parameters", []) - - for param in params: - param_str = str(param) - - # Extract component types from Query - query_match = re.findall(r'Query<[^>]*?([A-Z]\w+)', param_str) - for comp_name in query_match: - for comp_id in components: - comp_data = created_nodes.get(comp_id, {}) - if comp_data.get("name") == comp_name: - if not graph.graph.has_edge(system_id, comp_id): - graph.add_edge(system_id, comp_id, "queries_component") - connections += 1 - - # Extract resource types from Res - res_match = re.findall(r'Res(?:Mut)?<([A-Z]\w+)', param_str) - for res_name in res_match: - for res_id in resources: - res_data = created_nodes.get(res_id, {}) - if res_data.get("name") == res_name: - if not graph.graph.has_edge(system_id, res_id): - graph.add_edge(system_id, res_id, "uses_resource") - connections += 1 - - # Connect bundles to their component fields - for bundle_id in bundles: - bundle_data = created_nodes.get(bundle_id, {}) - m = bundle_data.get("metadata", {}) - fields = m.get("fields", []) - - for field in fields: - field_type = _extract_field_type(field) - for comp_id in components: - comp_data = created_nodes.get(comp_id, {}) - if comp_data.get("name") == field_type: - if not graph.graph.has_edge(bundle_id, comp_id): - graph.add_edge(bundle_id, comp_id, "bundles_component") - connections += 1 - - return connections - def _create_import_symbolic_nodes(graph, import_node_id, file_node_id, import_path, lang): """Create rich symbolic nodes for imports. Returns dict of new nodes to add.""" new_nodes = {} @@ -5007,52 +4926,95 @@ def _build_enhanced_result(chunk_text: str, meta: dict, score: float, graph: Opt def _get_result_graph_context(meta: dict, graph: Optional[LocalGraph] = None) -> dict: - """Ultra-compact graph context using LocalGraph helper methods.""" - if graph is None: + """Fixed graph context using enhanced graph methods.""" + if graph is None or graph.graph.number_of_nodes() == 0: return {} + context = {} file_path = meta.get('file', '') entity_name = meta.get('name', '') entity_type = meta.get('type', '') language = meta.get('language', '') - # Build node ID using graph's method - node_attrs = {"lang": language, "type": entity_type, "file": file_path, "name": entity_name} - node_id = graph.add_node(node_attrs, return_id=True) + # Try multiple node ID formats to find the right one + node_candidates = [] - if node_id not in graph.graph: - return {} + # Format 1: Standard node ID + if entity_name and file_path and language and entity_type: + node_candidates.append(f"{language}::{entity_type}::{file_path}::{entity_name}") - # Get all relationships in one call - relationships = graph.get_node_relationships(node_id) - context = {} + # Format 2: Method with class name + if entity_type in ['method', 'function'] and meta.get('class'): + class_name = meta.get('class') + node_candidates.append(f"{language}::{entity_type}::{file_path}::{class_name}.{entity_name}") - if relationships.get('cross_file'): - context['cross_file'] = relationships['cross_file'] - if relationships.get('same_file'): - context['same_file'] = relationships['same_file'] + # Format 3: Svelte component + if entity_type == 'component' and language == 'svelte': + node_candidates.append(f"svelte::component::{file_path}::{entity_name}") - # Build reverse dependencies from incoming relationships - if relationships.get('incoming'): - reverse_deps = [] - node_file = graph.graph.nodes[node_id].get('file', '') - for pred_id, edge_type, pred_data in relationships['incoming']: - pred_file = pred_data.get('file', '') - pred_name = pred_data.get('name', '') - pred_type = pred_data.get('type', '') + # Try to find the node using graph's enhanced search + target_node_id = None + for candidate in node_candidates: + if candidate in graph.graph.nodes: + target_node_id = candidate + break - if pred_file and pred_file != node_file: - reverse_deps.append(f"{edge_type}:{pred_type}:{pred_name}@{pred_file}") - else: - reverse_deps.append(f"{edge_type}:{pred_type}:{pred_name}") + # If still not found, try attribute-based search + if not target_node_id and entity_name and file_path: + matches = list(graph.find_nodes_by_attributes( + name=entity_name, + file=file_path, + type=entity_type if entity_type != 'code' else None + )) + if matches: + target_node_id = matches[0][0] # Take first match - if reverse_deps: - context['used_by'] = reverse_deps + if not target_node_id: + return context + + # Use the enhanced graph methods to get relationships + cross_file_out, cross_file_in = graph.get_cross_file_relationships(target_node_id, file_path) + + if cross_file_out: + context['cross_file'] = cross_file_out + if cross_file_in: + context['used_by'] = cross_file_in + + # Get same-file relationships + same_file_rels = [] + for neighbor_id in graph.graph.successors(target_node_id): + if neighbor_id in graph.graph.nodes: + neighbor_file = graph.graph.nodes[neighbor_id].get('file', '') + if neighbor_file == file_path: # Same file + edge_data = graph.graph.edges[target_node_id, neighbor_id] + rel_type = edge_data.get('type', 'related') + neighbor_name = graph.graph.nodes[neighbor_id].get('name', '') + neighbor_type = graph.graph.nodes[neighbor_id].get('type', '') + same_file_rels.append(f"{rel_type}:{neighbor_type}:{neighbor_name}") + + if same_file_rels: + context['same_file'] = same_file_rels + + # Get meaningful file entities (filter out noise) + meaningful_types = {'function', 'method', 'class', 'struct', 'interface', 'enum', 'module', 'component', 'trait', 'impl'} + file_entities = [] + + file_node_id = f"file::{file_path}" + if file_node_id in graph.graph.nodes: + for neighbor_id, neighbor_data in graph.neighbors(file_node_id): + neighbor_type = neighbor_data.get('type', '') + neighbor_name = neighbor_data.get('name', '') + + # Filter criteria + if (neighbor_type in meaningful_types and + neighbor_name and + len(neighbor_name) > 1 and # No single letters + neighbor_name != entity_name and + not neighbor_name.startswith(('c:', 'b:', 'p:', 'r:', 't:', 'v:'))): # No import aliases + file_entities.append(f"{neighbor_type}:{neighbor_name}") - # Get file entities - file_entities = graph.get_file_entities(file_path) if file_entities: - context['file_contents'] = file_entities + context['file_entities'] = file_entities return context