diff --git a/mcp_codebase.py b/mcp_codebase.py index 0714a37..40bcd0a 100644 --- a/mcp_codebase.py +++ b/mcp_codebase.py @@ -2248,19 +2248,40 @@ def build_indexes(): graph.add_node(node_id, **base_attrs, type=node_type) graph.add_edge(file_node_id, node_id, "contains") - # Handle imports if present + # ---------- IMPORTS ---------- 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") + # Find an existing module node with this name. + candidate_id = None + for node_id, data in created_nodes.items(): + if data.get("type") in {"module", "import"} and data.get("name") == imp: + candidate_id = node_id + break - # Handle function calls if present + if not candidate_id: + # If nothing exists, fall back to a generic import node. + candidate_id = f"import::{imp}" + if candidate_id not in graph.graph.nodes: + graph.add_node(candidate_id, type="Import", name=imp) + + graph.add_edge(file_node_id, candidate_id, "imports") + + # ---------- FUNCTION CALLS ---------- 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") + # Try to find a function/method node with this name. + target_id = None + for node_id, data in created_nodes.items(): + if data.get("type") in {"function", "method", "constructor"} and data.get("name") == call: + target_id = node_id + break + + if not target_id: + # Fallback to a generic call node. + target_id = f"call::{call}" + if target_id not in graph.graph.nodes: + graph.add_node(target_id, type="FunctionCall", name=call) + + graph.add_edge(node_id, target_id, "calls") + # ------------------------------------------------- # Step 2: Second-pass relationship building