From 9d73eb367953ead5cce2303044da5f3890660aca Mon Sep 17 00:00:00 2001
From: Stephen Adamson <popertots@gmail.com>
Date: Fri, 14 Nov 2025 22:11:59 +0000
Subject: [PATCH] Add helpers

---
 graph/graph.py  | 25 ++++++++++++++++++++++++-
 mcp_codebase.py |  9 +++++++++
 2 files changed, 33 insertions(+), 1 deletion(-)

diff --git a/graph/graph.py b/graph/graph.py
index 40c63ae..96fe75f 100644
--- a/graph/graph.py
+++ b/graph/graph.py
@@ -14,9 +14,32 @@ class LocalGraph:
         self.graph = nx.DiGraph()
 
     # --- Creation ---
-    def add_node(self, node_id: str, **attrs):
+    def add_node(self, node_id: str | dict, **attrs):
+        """
+        If a dict is passed instead of a plain string, build the node_id
+        as:  <lang>::<type>::<file>::<name>
+        """
+        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}"
         self.graph.add_node(node_id, **attrs)
 
+    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:
diff --git a/mcp_codebase.py b/mcp_codebase.py
index c671bf2..0dfd49d 100644
--- a/mcp_codebase.py
+++ b/mcp_codebase.py
@@ -3000,6 +3000,15 @@ def search_symbol(symbol: str, root_dir: Path, top_k: int = 20) -> List[Dict[str
         logger.warning(f"ripgrep failed: {e}")
     return matches
 
+
+def _extract_imports(source: str) -> list[str]:
+    """Return a list of fully‑qualified imports found in `source`."""
+    return [m.group(1).strip() for m in re.finditer(r'^\s*use\s+([^;]+);', source, re.MULTILINE)]
+
+def _extract_calls(source: str) -> list[str]:
+    """Return a list of called function names (bare name only)."""
+    return [m.group(1).split("::")[-1] for m in re.finditer(r'([A-Za-z_][A-Za-z0-9_:]*)\s*\(', source, re.MULTILINE)]
+
 # -----------------------------
 # Additional tool: read file lines with context
 # -----------------------------