diff --git a/mcp_codebase.py b/mcp_codebase.py index 0cfdb0c..f767ad8 100644 --- a/mcp_codebase.py +++ b/mcp_codebase.py @@ -2085,6 +2085,9 @@ def build_indexes(): graph = LocalGraph(root_path=str(CODEBASE_PATH)) graph.clear() + # Track all created nodes for second-pass relationships + created_nodes = {} + # Create graph nodes/edges from metadata with enhanced relationships for m in metadatas: fpath = m.get("file") @@ -2095,12 +2098,14 @@ def build_indexes(): component_name = m.get("component_name") # Skip empty/invalid nodes - if not name or node_type == "unknown": + if not name or node_type == "unknown" or not fpath: continue # Ensure file node exists file_node_id = f"file::{fpath}" - graph.add_node(file_node_id, type="File", path=fpath, lang=lang) + if file_node_id not in graph.graph.nodes: + graph.add_node(file_node_id, type="File", path=fpath, lang=lang) + created_nodes[file_node_id] = {"type": "File", "path": fpath} # Create node ID with language node_id = f"{lang}::{node_type}::{fpath}::{name}" @@ -2113,19 +2118,30 @@ def build_indexes(): "lang": lang # Always include language } + # Store for second-pass processing + created_nodes[node_id] = base_attrs.copy() + created_nodes[node_id]["type"] = node_type + # Create appropriate node based on type if node_type in ["function", "method", "constructor"]: # Function-like entities if node_type == "method" and "class" in m: # Method belongs to a class - use compound ID - node_id = f"{lang}::{node_type}::{fpath}::{m['class']}.{name}" + class_name = m['class'] + node_id = f"{lang}::{node_type}::{fpath}::{class_name}.{name}" + + # Update created_nodes with new ID + created_nodes[node_id] = base_attrs.copy() + created_nodes[node_id]["type"] = node_type + created_nodes[node_id]["class_name"] = class_name + graph.add_node(node_id, **base_attrs, type=node_type, - class_name=m.get("class"), + class_name=class_name, return_type=m.get("return_type"), parameters=m.get("parameters", [])) # Connect method to its class - class_node_id = f"{lang}::class::{fpath}::{m['class']}" + class_node_id = f"{lang}::class::{fpath}::{class_name}" graph.add_edge(class_node_id, node_id, "contains") # Connect return type if available @@ -2153,6 +2169,7 @@ def build_indexes(): graph.add_node(node_id, **base_attrs, type=node_type, return_type=m.get("return_type"), parameters=m.get("parameters", [])) + # Connect function to file graph.add_edge(file_node_id, node_id, "contains") # Connect return type if available @@ -2181,27 +2198,17 @@ def build_indexes(): graph.add_node(node_id, **base_attrs, type=node_type, extends=m.get("extends"), implements=m.get("implements", []), - fields=m.get("fields", []), # Store fields as metadata, not separate nodes - variants=m.get("variants", [])) # For enums + fields=m.get("fields", []), + variants=m.get("variants", [])) + # Connect type to file graph.add_edge(file_node_id, node_id, "contains") - # Inheritance relationships - extends_class = m.get("extends") - if extends_class: - for parent_id, parent_data in graph.find_nodes(name=extends_class): - graph.add_edge(node_id, parent_id, "extends") - - # Interface implementation - implements_interfaces = m.get("implements", []) - for interface_name in implements_interfaces: - for interface_id, interface_data in graph.find_nodes(name=interface_name): - graph.add_edge(node_id, interface_id, "implements") - - # Bevy-specific: Plugin implementation detection - if node_type == "struct" and name.endswith("Plugin"): - # FIX: Use add_node with attributes instead of modifying existing node - graph.add_node(node_id, **base_attrs, type=node_type, is_bevy_plugin=True) + # Store additional attributes for second-pass + created_nodes[node_id]["extends"] = m.get("extends") + created_nodes[node_id]["implements"] = m.get("implements", []) + created_nodes[node_id]["fields"] = m.get("fields", []) + created_nodes[node_id]["variants"] = m.get("variants", []) elif node_type == "impl": # Implementation blocks @@ -2213,21 +2220,15 @@ def build_indexes(): traits=m.get("traits", []), methods=m.get("methods", [])) + # Connect impl to file graph.add_edge(file_node_id, node_id, "contains") - # Connect to target type - for target_id, target_data in graph.find_nodes(name=target): - graph.add_edge(node_id, target_id, "implements_for") - - # Connect to implemented traits - for trait_name in m.get("traits", []): - for trait_id, trait_data in graph.find_nodes(name=trait_name): - graph.add_edge(node_id, trait_id, "implements") - - # Bevy-specific: Plugin detection - if target.endswith("Plugin"): - # FIX: Use add_node with attributes instead of modifying existing node - graph.add_node(node_id, **base_attrs, type=node_type, is_bevy_plugin_impl=True) + # Store for second-pass + created_nodes[node_id] = base_attrs.copy() + created_nodes[node_id]["type"] = node_type + created_nodes[node_id]["target"] = target + created_nodes[node_id]["traits"] = m.get("traits", []) + created_nodes[node_id]["methods"] = m.get("methods", []) elif node_type == "trait": # Traits @@ -2235,6 +2236,8 @@ def build_indexes(): methods=m.get("methods", [])) graph.add_edge(file_node_id, node_id, "contains") + created_nodes[node_id]["methods"] = m.get("methods", []) + elif node_type == "module": # Modules graph.add_node(node_id, **base_attrs, type=node_type) @@ -2247,104 +2250,30 @@ def build_indexes(): # Handle imports if present for imp in m.get("imports", []): + # Create import node if it doesn't exist + if imp not in graph.graph.nodes: + graph.add_node(imp, type="Import", name=imp) graph.add_edge(file_node_id, imp, "imports") # Handle function calls if present for call in m.get("calls", []): + # Create call target node if it doesn't exist + if call not in graph.graph.nodes: + graph.add_node(call, type="FunctionCall", name=call) graph.add_edge(node_id, call, "calls") - # Rust-specific relationships - FIX: Use graph.graph.nodes - if lang == "rust": - if node_type == "struct": - # Connect struct fields as metadata, not separate nodes - fields = m.get("fields", []) - for field in fields: - # Store field info as node attribute instead of creating separate nodes - if "field_types" not in graph.graph.nodes[node_id]: - graph.graph.nodes[node_id]["field_types"] = [] - graph.graph.nodes[node_id]["field_types"].append(field) + # ------------------------------------------------- + # Step 2: Second-pass relationship building + # ------------------------------------------------- + logger.info("Building second-pass relationships...") - elif node_type == "enum": - # Connect enum variants as metadata - variants = m.get("variants", []) - for variant in variants: - if "variants" not in graph.graph.nodes[node_id]: - graph.graph.nodes[node_id]["variants"] = [] - graph.graph.nodes[node_id]["variants"].append(variant) + # Build type relationships + type_connections = build_type_relationships(graph, created_nodes) + logger.info(f"Added {type_connections} type relationships") - elif node_type == "function": - # Connect function parameters to their types - parameters = m.get("parameters", []) - for param in parameters: - # Extract type from "Type param_name" format - param_parts = param.split() - if len(param_parts) >= 2: - param_type = param_parts[0] - if param_type not in ['self', '&self', '&mut self', 'mut']: - for type_id, type_data in graph.find_nodes(name=param_type): - graph.add_edge(node_id, type_id, "uses_parameter") - - # Bevy-specific relationship detection - FIX: Use graph.graph.nodes - if lang == "rust" and node_type == "function": - # Detect Bevy system functions - function_name = name.lower() - if any(keyword in function_name for keyword in ['system', 'plugin', 'build', 'update']): - graph.graph.nodes[node_id]["is_bevy_system"] = True - - # Detect query parameters for system dependencies - parameters = m.get("parameters", []) - for param in parameters: - if 'Query' in param: - graph.graph.nodes[node_id]["has_bevy_query"] = True - if 'Res<' in param or 'ResMut<' in param: - graph.graph.nodes[node_id]["has_bevy_resource"] = True - if 'Commands' in param: - graph.graph.nodes[node_id]["has_bevy_commands"] = True - - # Second pass for Bevy-specific relationships - FIX: Use graph.graph.nodes - bevy_connections = 0 - - for node_id, node_data in graph.graph.nodes(data=True): - if node_data.get('lang') == 'rust': - # Connect Plugin impls to their build functions - if node_data.get('type') == 'impl' and node_data.get('is_bevy_plugin_impl'): - plugin_struct = node_data.get('target') - if plugin_struct: - # Find the build method in this impl - for method_id, method_data in graph.find_nodes(type="method", class_name=plugin_struct): - if method_data.get('name') == 'build': - graph.add_edge(node_id, method_id, "defines_plugin") - bevy_connections += 1 - - # Connect systems to their plugins - elif node_data.get('is_bevy_system'): - # Try to find which plugin this system belongs to - file_path = node_data.get('file', '') - if file_path: - # Look for plugins in the same file - for plugin_id, plugin_data in graph.find_nodes(type="struct", file=file_path): - if plugin_data.get('name', '').endswith('Plugin'): - graph.add_edge(plugin_id, node_id, "contains_system") - bevy_connections += 1 - - if bevy_connections > 0: - logger.info(f"Added {bevy_connections} Bevy-specific relationships") - - logger.info("Building method parameter type relationships...") - param_connections = 0 - for node_id, node_data in graph.graph.nodes(data=True): - if node_data.get('type') in ['method', 'function', 'constructor']: - parameters = node_data.get('parameters', []) - for param in parameters: - # Extract type from parameter string "TypeName paramName" - param_type = param.split()[0] if param and ' ' in param else None - if param_type and param_type not in ['void', 'int', 'String', 'boolean', 'long', 'double', 'float']: - # Try to find this type in the graph - for type_id, type_data in graph.find_nodes(name=param_type): - graph.add_edge(node_id, type_id, "uses_parameter") - param_connections += 1 - - logger.info(f"Added {param_connections} parameter type relationships") + # Build Bevy-specific relationships + bevy_connections = build_bevy_relationships(graph, created_nodes) + logger.info(f"Added {bevy_connections} Bevy-specific relationships") graph.save() graph.to_json() @@ -2419,6 +2348,206 @@ def build_indexes(): elapsed = index_build_time - start_time logger.info(f"Index build complete with contextual weighting applied. Total time: {elapsed:.1f}s") +def build_type_relationships(graph, created_nodes): + """Build type relationships between nodes (parameter types, return types, etc.)""" + connections = 0 + + for node_id, node_data in created_nodes.items(): + node_type = node_data.get("type") + + # Handle function/method parameter and return types + if node_type in ["function", "method", "constructor"]: + connections += _connect_function_types(graph, node_id, node_data) + + # Handle class/struct inheritance and implementation + elif node_type in ["class", "struct", "interface"]: + connections += _connect_class_relationships(graph, node_id, node_data) + + # Handle impl blocks + elif node_type == "impl": + connections += _connect_impl_relationships(graph, node_id, node_data) + + return connections + +def _connect_function_types(graph, node_id, node_data): + """Connect function nodes to their parameter and return types""" + connections = 0 + + # Connect return type + return_type = node_data.get("return_type") + if return_type and _is_meaningful_type(return_type): + return_type_clean = _clean_type_name(return_type) + for type_id, type_data in graph.find_nodes(name=return_type_clean): + graph.add_edge(node_id, type_id, "returns") + connections += 1 + + # Connect parameter types + parameters = node_data.get("parameters", []) + for param in parameters: + param_type = _extract_parameter_type(param) + if param_type and _is_meaningful_type(param_type): + param_type_clean = _clean_type_name(param_type) + for type_id, type_data in graph.find_nodes(name=param_type_clean): + graph.add_edge(node_id, type_id, "uses_parameter") + connections += 1 + + return connections + +def _connect_class_relationships(graph, node_id, node_data): + """Connect class/struct nodes to parent classes and interfaces""" + connections = 0 + + # Inheritance + extends = node_data.get("extends") + if extends and _is_meaningful_type(extends): + for parent_id, parent_data in graph.find_nodes(name=extends): + graph.add_edge(node_id, parent_id, "extends") + connections += 1 + + # Interface implementation + implements = node_data.get("implements", []) + for interface in implements: + if _is_meaningful_type(interface): + for interface_id, interface_data in graph.find_nodes(name=interface): + graph.add_edge(node_id, interface_id, "implements") + connections += 1 + + return connections + +def _connect_impl_relationships(graph, node_id, node_data): + """Connect impl blocks to their targets and traits""" + connections = 0 + + target = node_data.get("target") + if target and _is_meaningful_type(target): + # Connect to target type + for target_id, target_data in graph.find_nodes(name=target): + graph.add_edge(node_id, target_id, "implements_for") + connections += 1 + + # Connect to implemented traits + traits = node_data.get("traits", []) + for trait in traits: + if _is_meaningful_type(trait): + for trait_id, trait_data in graph.find_nodes(name=trait): + graph.add_edge(node_id, trait_id, "implements") + connections += 1 + + return connections + +def build_bevy_relationships(graph, created_nodes): + """Build Bevy-specific relationships""" + connections = 0 + + # Find all plugin structs and impls + plugin_structs = {} + plugin_impls = {} + + for node_id, node_data in created_nodes.items(): + node_type = node_data.get("type") + name = node_data.get("name", "") + + # Collect plugin structs + if node_type in ["struct", "class"] and name.endswith("Plugin"): + plugin_structs[name] = node_id + + # Collect plugin impls + elif node_type == "impl" and node_data.get("target", "").endswith("Plugin"): + plugin_impls[node_data["target"]] = node_id + + # Connect plugin impls to their structs + for plugin_name, impl_id in plugin_impls.items(): + if plugin_name in plugin_structs: + struct_id = plugin_structs[plugin_name] + graph.add_edge(impl_id, struct_id, "implements_plugin") + connections += 1 + + # Find build method in this impl + impl_methods = created_nodes[impl_id].get("methods", []) + for method_name in impl_methods: + if method_name == "build": + # Find the method node + method_id = f"{created_nodes[impl_id]['lang']}::method::{created_nodes[impl_id]['file']}::{plugin_name}.{method_name}" + if method_id in created_nodes: + graph.add_edge(impl_id, method_id, "defines_plugin") + connections += 1 + + # Connect systems to plugins (simplified - same file assumption) + for node_id, node_data in created_nodes.items(): + if node_data.get("type") in ["function", "method"]: + # Simple heuristic for Bevy systems + function_name = node_data.get("name", "").lower() + if any(keyword in function_name for keyword in ['system', 'update', 'setup']): + file_path = node_data.get("file") + # Look for plugins in the same file + for plugin_id, plugin_data in created_nodes.items(): + if (plugin_data.get("file") == file_path and + plugin_data.get("type") in ["struct", "class"] and + plugin_data.get("name", "").endswith("Plugin")): + graph.add_edge(plugin_id, node_id, "contains_system") + connections += 1 + + return connections + +def _extract_parameter_type(param): + """Extract type from parameter string, handling multiple formats""" + if not param: + return None + + # Handle "type: name" format (TypeScript/Python) + if ':' in param: + parts = param.split(':') + if len(parts) >= 2: + return parts[1].strip() + + # Handle "Type name" format (Rust/Java/C++) + if ' ' in param: + parts = param.split() + if len(parts) >= 2: + return parts[0].strip() + + # Handle complex types with generics + if '<' in param and '>' in param: + # Extract the base type before generics + base_type = param.split('<')[0].strip() + return base_type + + return param.strip() + +def _clean_type_name(type_name): + """Clean type names by removing common modifiers""" + if not type_name: + return "" + + # Remove common modifiers + modifiers = ['&', 'mut ', 'const ', 'static ', 'pub ', 'private ', 'protected '] + cleaned = type_name + for mod in modifiers: + cleaned = cleaned.replace(mod, '') + + # Remove trailing references/pointers + cleaned = cleaned.rstrip('*&') + + return cleaned.strip() + +def _is_meaningful_type(type_name): + """Check if a type name is meaningful (not primitive, void, etc.)""" + if not type_name: + return False + + meaningless_types = [ + 'void', 'int', 'i32', 'i64', 'f32', 'f64', 'bool', 'str', 'String', + 'char', 'u8', 'u16', 'u32', 'u64', 'usize', 'isize', + 'self', '&self', '&mut self', 'mut self', 'Self', + 'Result', 'Option', 'Vec', 'String' + ] + + cleaned = _clean_type_name(type_name) + return (cleaned and + cleaned not in meaningless_types and + not cleaned.startswith('&') and + len(cleaned) > 1) + # ----------------------------- # Build / load indexes with contextual weighting and batch embedding # -----------------------------