context-mcp-server 1.0.1

package/codegraph/extractors/ast_extractor.py

@@ -0,0 +1,222 @@
+"""
+ast_extractor.py — extract nodes from code files using tree-sitter AST.
+
+Falls back to regex if tree-sitter grammars aren't installed.
+Each node: { id, name, type, file, line, docstring?, calls?, imports? }
+"""
+
+import re
+from pathlib import Path
+
+
+# ── tree-sitter setup (optional — graceful fallback) ─────────────────────────
+
+def _try_load_ts():
+    try:
+        import tree_sitter_python as tspython
+        import tree_sitter_javascript as tsjavascript
+        from tree_sitter import Language, Parser
+        PY_LANG = Language(tspython.language())
+        JS_LANG = Language(tsjavascript.language())
+        return {"python": (PY_LANG, Parser(PY_LANG)), "javascript": (JS_LANG, Parser(JS_LANG))}
+    except ImportError:
+        return {}
+
+_TS_PARSERS = _try_load_ts()
+
+
+# ── tree-sitter queries ───────────────────────────────────────────────────────
+
+_PY_CLASS_QUERY = """
+(class_definition name: (identifier) @name) @class
+"""
+
+_PY_FUNC_QUERY = """
+(function_definition name: (identifier) @name) @func
+"""
+
+_JS_CLASS_QUERY = """
+(class_declaration name: (identifier) @name) @class
+"""
+
+_JS_FUNC_QUERY = """
+[
+  (function_declaration name: (identifier) @name)
+  (method_definition name: (property_identifier) @name)
+] @func
+"""
+
+
+def _ts_extract(source: bytes, lang_key: str, rel_path: str) -> list:
+    parsers = _TS_PARSERS
+    if lang_key not in parsers:
+        return []
+    lang, parser = parsers[lang_key]
+    tree = parser.parse(source)
+
+    nodes = []
+    lines = source.decode("utf-8", errors="replace").splitlines()
+
+    def _node(kind, name, line):
+        return {
+            "id": f"{rel_path}::{kind}::{name}",
+            "name": name,
+            "type": kind,
+            "file": rel_path,
+            "line": line,
+        }
+
+    def _iter_captures(query, root):
+        """Yield (capture_name, tree_node) pairs; compatible with tree-sitter >=0.20."""
+        try:
+            # tree-sitter >= 0.22: matches() returns list of (pattern_idx, {name: [Node]})
+            for _pat_idx, caps in query.matches(root):
+                for cap_name, cap_nodes in caps.items():
+                    for n in (cap_nodes if isinstance(cap_nodes, list) else [cap_nodes]):
+                        yield cap_name, n
+        except Exception:
+            pass
+
+    # Classes
+    try:
+        query = lang.query(_PY_CLASS_QUERY if lang_key == "python" else _JS_CLASS_QUERY)
+        for cap_name, node in _iter_captures(query, tree.root_node):
+            if cap_name == "name" and node.type == "identifier":
+                nodes.append(_node("class", node.text.decode(), node.start_point[0] + 1))
+    except Exception:
+        pass
+
+    # Functions
+    try:
+        query = lang.query(_PY_FUNC_QUERY if lang_key == "python" else _JS_FUNC_QUERY)
+        for cap_name, node in _iter_captures(query, tree.root_node):
+            if cap_name == "name" and node.type in ("identifier", "property_identifier"):
+                nodes.append(_node("function", node.text.decode(), node.start_point[0] + 1))
+    except Exception:
+        pass
+
+    return nodes
+
+
+# ── Regex fallback ────────────────────────────────────────────────────────────
+
+_PATTERNS = {
+    "python": {
+        "class":    re.compile(r"^class\s+(\w+)", re.MULTILINE),
+        "function": re.compile(r"^def\s+(\w+)", re.MULTILINE),
+        "import":   re.compile(r"^(?:import|from)\s+([\w.]+)", re.MULTILINE),
+    },
+    "javascript": {
+        "class":    re.compile(r"\bclass\s+(\w+)", re.MULTILINE),
+        "function": re.compile(r"\bfunction\s+(\w+)", re.MULTILINE),
+        "import":   re.compile(r"^import\s+.*?from\s+['\"](.+?)['\"]", re.MULTILINE),
+    },
+    "go": {
+        "function": re.compile(r"^func\s+(?:\(\w+\s+\*?\w+\)\s+)?(\w+)", re.MULTILINE),
+        "struct":   re.compile(r"^type\s+(\w+)\s+struct", re.MULTILINE),
+        "import":   re.compile(r'"([\w./]+)"', re.MULTILINE),
+    },
+    "rust": {
+        "function": re.compile(r"^(?:pub\s+)?fn\s+(\w+)", re.MULTILINE),
+        "struct":   re.compile(r"^(?:pub\s+)?struct\s+(\w+)", re.MULTILINE),
+        "import":   re.compile(r"^use\s+([\w:]+)", re.MULTILINE),
+    },
+    "java": {
+        "class":    re.compile(r"\bclass\s+(\w+)", re.MULTILINE),
+        "function": re.compile(r"(?:public|private|protected|static|\s)+[\w<>\[\]]+\s+(\w+)\s*\(", re.MULTILINE),
+        "import":   re.compile(r"^import\s+([\w.]+);", re.MULTILINE),
+    },
+    "ruby": {
+        "class":    re.compile(r"^class\s+(\w+)", re.MULTILINE),
+        "function": re.compile(r"^\s*def\s+(\w+)", re.MULTILINE),
+    },
+    "sql": {
+        "table":    re.compile(r"CREATE\s+TABLE\s+(?:IF\s+NOT\s+EXISTS\s+)?[`\"]?(\w+)[`\"]?", re.IGNORECASE),
+        "index":    re.compile(r"CREATE\s+(?:UNIQUE\s+)?INDEX\s+\w+\s+ON\s+[`\"]?(\w+)[`\"]?", re.IGNORECASE),
+    },
+    # config files: no node extraction — file node created by scanner
+}
+
+
+def _ext_to_lang(ext: str) -> str:
+    return {
+        ".py": "python", ".pyw": "python",
+        ".js": "javascript", ".mjs": "javascript", ".cjs": "javascript",
+        ".jsx": "javascript", ".ts": "javascript", ".tsx": "javascript",
+        ".go": "go", ".rs": "rust",
+        ".java": "java", ".rb": "ruby",
+        ".sql": "sql",
+        ".yaml": "config", ".yml": "config", ".toml": "config",
+        ".env": "config", ".ini": "config", ".cfg": "config",
+    }.get(ext, "")
+
+
+def _regex_extract(text: str, lang: str, rel_path: str) -> list:
+    patterns = _PATTERNS.get(lang, {})
+    children = []
+    for kind, pat in patterns.items():
+        if kind == "import":
+            continue
+        for m in pat.finditer(text):
+            line = text[:m.start()].count("\n") + 1
+            children.append({
+                "id":   f"{rel_path}::{kind}::{m.group(1)}",
+                "name": m.group(1),
+                "type": kind,
+                "file": rel_path,
+                "line": line,
+            })
+
+    imp_pat = patterns.get("import")
+    imports = [m.group(1) for m in imp_pat.finditer(text)][:30] if imp_pat else []
+
+    return _wrap_in_module(rel_path, children, imports)
+
+
+# ── Module wrapper ───────────────────────────────────────────────────────────
+
+def _wrap_in_module(rel_path: str, children: list, imports: list) -> list:
+    """
+    Create a module node for the file, link all child nodes to it via defined_in.
+    The module node carries the imports so builder can create file-to-file edges.
+    """
+    from pathlib import Path as _Path
+    stem = _Path(rel_path).stem
+    mod_id = f"{rel_path}::module::{stem}"
+    module = {
+        "id":      mod_id,
+        "name":    stem,
+        "type":    "module",
+        "file":    rel_path,
+        "line":    1,
+        "imports": imports,
+    }
+    for child in children:
+        child["relations"] = [{"id": mod_id, "relation": "defined-in", "confidence": "EXTRACTED"}]
+    return [module] + children
+
+
+# ── Public API ────────────────────────────────────────────────────────────────
+
+def extract(abs_path: str, rel_path: str) -> list:
+    """Extract nodes from a code/sql/config file. Returns list of node dicts."""
+    ext = Path(abs_path).suffix.lower()
+    lang = _ext_to_lang(ext)
+    if not lang:
+        return []
+
+    try:
+        raw = open(abs_path, "rb").read()
+        text = raw.decode("utf-8", errors="replace")
+    except OSError:
+        return []
+
+    # tree-sitter for Python and JS/TS if available
+    if lang in ("python", "javascript") and lang in _TS_PARSERS:
+        children = _ts_extract(raw, lang, rel_path)
+        if children:
+            imp_pat = _PATTERNS.get(lang, {}).get("import")
+            imports = [m.group(1) for m in imp_pat.finditer(text)][:30] if imp_pat else []
+            return _wrap_in_module(rel_path, children, imports)
+
+    return _regex_extract(text, lang, rel_path)

package/codegraph/extractors/audio_extractor.py

@@ -0,0 +1,8 @@
+"""
+audio_extractor.py — audio files are not supported without faster-whisper.
+Stub kept so imports don't break.
+"""
+
+
+def transcribe(_path: str) -> str:
+    return ""

package/codegraph/extractors/doc_extractor.py

@@ -0,0 +1,34 @@
+"""
+doc_extractor.py — extract plain text from doc and PDF files.
+PDF extraction uses pymupdf if installed; falls back to label-only otherwise.
+"""
+
+from pathlib import Path
+
+
+def extract_text(path: str) -> str:
+    """Return text content of a doc/PDF file. Truncated at DOC_MAX_CHARS."""
+    from ..config import DOC_MAX_CHARS
+    if path.lower().endswith(".pdf"):
+        return _extract_pdf(path, DOC_MAX_CHARS)
+    try:
+        return Path(path).read_text(encoding="utf-8", errors="replace")[:DOC_MAX_CHARS]
+    except OSError:
+        return ""
+
+
+def _extract_pdf(path: str, max_chars: int) -> str:
+    try:
+        import pymupdf  # optional dep
+        doc = pymupdf.open(path)
+        parts = []
+        for page in doc:
+            parts.append(page.get_text())
+            if sum(len(p) for p in parts) >= max_chars:
+                break
+        doc.close()
+        return "".join(parts)[:max_chars]
+    except ImportError:
+        return f"[PDF: {Path(path).name} — install pymupdf to extract text]"
+    except Exception:
+        return ""

package/codegraph/extractors/image_extractor.py

@@ -0,0 +1,26 @@
+"""
+image_extractor.py — encode images as base64 for AI vision.
+No external deps — stdlib only.
+"""
+
+import base64
+import mimetypes
+from pathlib import Path
+
+
+def extract_image_b64(path: str) -> dict | None:
+    """Return {"data": base64_str, "media_type": "image/png"} or None on failure."""
+    try:
+        data = Path(path).read_bytes()
+        media_type = mimetypes.guess_type(path)[0] or "image/png"
+        return {"data": base64.b64encode(data).decode(), "media_type": media_type}
+    except OSError:
+        return None
+
+
+def extract_svg_text(path: str) -> str:
+    """Return SVG file as plain text (SVGs are XML — readable as-is)."""
+    try:
+        return Path(path).read_text(encoding="utf-8", errors="replace")[:4000]
+    except OSError:
+        return ""

package/codegraph/graph/builder.py

@@ -0,0 +1,145 @@
+"""
+graph/builder.py — build a NetworkX directed graph from extracted nodes.
+
+Node attributes: id, name, type, file, line, community
+Edge attributes: relation, confidence (EXTRACTED | INFERRED | AMBIGUOUS)
+"""
+
+import json
+from pathlib import Path
+
+try:
+    import networkx as nx
+    _HAS_NX = True
+except ImportError:
+    _HAS_NX = False
+
+
+def build(all_nodes: list[dict]) -> "nx.DiGraph | dict":
+    """
+    Build graph from flat node list. Returns nx.DiGraph or plain dict fallback.
+    Edges are created from node.imports[] and node.calls[] fields.
+    """
+    if not _HAS_NX:
+        return _dict_graph(all_nodes)
+
+    G = nx.DiGraph()
+
+    node_by_name: dict[str, str] = {}    # name -> id
+    module_by_file: dict[str, str] = {}  # rel_path -> module node id
+
+    for node in all_nodes:
+        nid = node.get("id", "")
+        if not nid:
+            continue
+        G.add_node(nid, **{k: v for k, v in node.items() if k not in ("imports", "calls", "relations")})
+        node_by_name[node.get("name", "")] = nid
+        if node.get("type") == "module":
+            module_by_file[node.get("file", "")] = nid
+
+    # Build file-path lookup from module nodes
+    file_node: dict[str, str] = {}
+    for rel_path, mod_id in module_by_file.items():
+        p = rel_path.replace("\\", "/")
+        stem = p.split("/")[-1].split(".")[0]
+        base = p.split("/")[-1]
+        for key in (stem, base, p):
+            file_node.setdefault(key, mod_id)
+
+    # defined-in edges: child nodes → their module
+    for node in all_nodes:
+        nid = node.get("id", "")
+        for rel in node.get("relations", []):
+            target_id = rel.get("id") or node_by_name.get(rel.get("name", ""))
+            if target_id and target_id != nid:
+                G.add_edge(nid, target_id,
+                           relation=rel.get("relation", "relates-to"),
+                           confidence=rel.get("confidence", "EXTRACTED"))
+
+    # Import edges: module → module
+    seen_edges: set[tuple] = set()
+    for node in all_nodes:
+        if node.get("type") != "module":
+            continue
+        src_id = node.get("id", "")
+        for imp in node.get("imports", []):
+            clean = imp.lstrip(".")
+            parts = clean.replace("\\", "/").split("/")
+            last  = parts[-1]
+            stem  = last.split(".")[0]
+            for c in (clean, last, stem):
+                if not c:
+                    continue
+                target = file_node.get(c) or node_by_name.get(c)
+                if target and target != src_id:
+                    key = (src_id, target)
+                    if key not in seen_edges:
+                        seen_edges.add(key)
+                        G.add_edge(src_id, target, relation="imports", confidence="EXTRACTED")
+                    break
+
+    # Edges from explicit relations (concept nodes from LLM)
+    for node in all_nodes:
+        nid = node.get("id", "")
+        for rel in node.get("relations", []):
+            target_id = rel.get("id") or node_by_name.get(rel.get("name", ""))
+            if target_id and target_id != nid:
+                G.add_edge(nid, target_id,
+                           relation=rel.get("relation", "relates-to"),
+                           confidence=rel.get("confidence", "INFERRED"))
+
+    return G
+
+
+def _dict_graph(all_nodes: list[dict]) -> dict:
+    """Fallback when networkx not installed."""
+    nodes = []
+    edges = []
+    seen = set()
+    for node in all_nodes:
+        nid = node.get("id", "")
+        if nid in seen:
+            continue
+        seen.add(nid)
+        nodes.append({k: v for k, v in node.items() if k not in ("imports", "calls", "relations")})
+        for imp in node.get("imports", []):
+            edges.append({"from": nid, "to": imp, "relation": "imports", "confidence": "EXTRACTED"})
+    return {"nodes": nodes, "edges": edges, "communities": [], "god_nodes": []}
+
+
+def to_json_dict(G) -> dict:
+    """Serialize graph to the graph.json schema."""
+    if isinstance(G, dict):
+        return G  # fallback path
+
+    nodes = [{"id": nid, **data} for nid, data in G.nodes(data=True)]
+    edges = [{"from": u, "to": v, **data} for u, v, data in G.edges(data=True)]
+
+    # God nodes = highest degree
+    degrees = sorted(G.degree(), key=lambda x: x[1], reverse=True)
+    god_nodes = [n for n, d in degrees[:5] if d > 2]
+
+    return {
+        "nodes": nodes,
+        "edges": edges,
+        "communities": G.graph.get("communities", []),
+        "god_nodes": god_nodes,
+        "generated_at": __import__("datetime").datetime.utcnow().isoformat() + "Z",
+    }
+
+
+def save_graph(project_root: str, graph_dict: dict) -> str:
+    out = Path(project_root) / "codegraph-cache" / "graph.json"
+    out.parent.mkdir(parents=True, exist_ok=True)
+    out.write_text(json.dumps(graph_dict, indent=2), encoding="utf-8")
+    return str(out)
+
+
+def load_graph(project_root: str) -> dict | None:
+    p = Path(project_root) / "codegraph-cache" / "graph.json"
+    if not p.exists():
+        return None
+    try:
+        return json.loads(p.read_text(encoding="utf-8"))
+    except Exception:
+        return None

package/codegraph/graph/clustering.py

@@ -0,0 +1,40 @@
+"""
+graph/clustering.py — community detection using NetworkX connected-components.
+No external deps beyond networkx (already required).
+"""
+
+
+def detect_communities(G) -> list[dict]:
+    """Assign community IDs to graph nodes. Returns list of community dicts."""
+    try:
+        import networkx as nx
+    except ImportError:
+        return []
+
+    if G.number_of_nodes() == 0:
+        return []
+
+    undirected = G.to_undirected()
+    communities = []
+    for comm_id, component in enumerate(nx.connected_components(undirected)):
+        member_ids = list(component)
+        label = _community_label(G, member_ids)
+        communities.append({"id": comm_id, "label": label, "members": member_ids})
+        for nid in member_ids:
+            if G.has_node(nid):
+                G.nodes[nid]["community"] = comm_id
+
+    G.graph["communities"] = communities
+    return communities
+
+
+def _community_label(G, member_ids: list) -> str:
+    files = []
+    for nid in member_ids:
+        if G.has_node(nid):
+            f = G.nodes[nid].get("file", "")
+            if f:
+                files.append(f.split("/")[0])
+    if not files:
+        return "misc"
+    return max(set(files), key=files.count)