Initial commit: Pegas Memory (PG + Qdrant + Memgraph, MCP server, 14 tools)

pegas_memory/code_graph.py

@@ -0,0 +1,164 @@
+"""Граф зависимостей кода: calls/imports/inherits через tree-sitter.
+
+Строит ориентированный граф зависимостей между символами кода:
+  - calls:    функция A вызывает функцию B
+  - imports:  файл импортирует модуль (src = имя файла без расширения)
+  - inherits: класс B наследует класс A
+
+Извлекает из tree-sitter AST и сохраняет в pm_code_dependency (PostgreSQL).
+"""
+from __future__ import annotations
+
+from .code_symbols import _parser
+
+_FUNCTION_TYPES = ("function_definition", "function_item", "function_declaration", "method_definition")
+_CLASS_TYPES = ("class_definition", "class_declaration")
+
+
+def _sym_name(node, sb: bytes) -> str | None:
+    n = node.child_by_field_name("name")
+    return sb[n.start_byte:n.end_byte].decode("utf-8", "replace") if n else None
+
+
+def _call_target(node, sb: bytes) -> str | None:
+    """Имя вызываемой функции из call-узла."""
+    fn = node.child_by_field_name("function")
+    if fn is None:
+        return None
+    if fn.type == "identifier":
+        return sb[fn.start_byte:fn.end_byte].decode("utf-8", "replace")
+    if fn.type == "attribute":
+        parts = [sb[c.start_byte:c.end_byte].decode("utf-8", "replace")
+                 for c in fn.children if c.type == "identifier"]
+        return parts[-1] if parts else None
+    return None
+
+
+def _collect_calls(node, owner, sb, add_fn):
+    """Рекурсивно собрать вызовы внутри тела функции/метода."""
+    if node.type == "call":
+        t = _call_target(node, sb)
+        if t:
+            add_fn(owner, t, "calls")
+    for ch in node.children:
+        _collect_calls(ch, owner, sb, add_fn)
+
+
+def _collect_inherits(node, owner, sb, add_fn):
+    """Найти базовые классы (superclasses/heritage_clause)."""
+    for ch in node.children:
+        if ch.type in ("argument_list", "superclasses", "heritage_clause", "base_list"):
+            for sub in _iter(ch):
+                if sub.type == "identifier":
+                    add_fn(owner, sb[sub.start_byte:sub.end_byte].decode("utf-8", "replace"), "inherits")
+                elif sub.type == "attribute":
+                    parts = [sb[p.start_byte:p.end_byte].decode("utf-8", "replace")
+                             for p in sub.children if p.type == "identifier"]
+                    if parts:
+                        add_fn(owner, parts[-1], "inherits")
+
+
+def _collect_imports(root, module_name, sb, add_fn):
+    """Собрать импорты (import X / from Y import Z)."""
+    for node in _iter(root):
+        if node.type == "import_statement":
+            for ch in node.children:
+                if ch.type in ("dotted_name", "identifier"):
+                    add_fn(module_name, sb[ch.start_byte:ch.end_byte].decode("utf-8", "replace"), "imports")
+                    break
+        elif node.type == "import_from_statement":
+            mod = node.child_by_field_name("module_name")
+            if mod is not None:
+                target = sb[mod.start_byte:mod.end_byte].decode("utf-8", "replace")
+                # отфильтровать артефакты: "." (relative import), "..", пустые
+                if target not in (".", "..", ""):
+                    add_fn(module_name, target, "imports")
+
+
+def _iter(node):
+    yield node
+    for ch in node.children:
+        yield from _iter(ch)
+
+
+def extract_dependencies(source: str, language: str, repo: str = "", path: str = "") -> list[dict]:
+    """Вернуть зависимости [{repo,path,language,src,dst,kind}] из исходника."""
+    parser, _ = _parser(language)
+    sb = source.encode("utf-8")
+    tree = parser.parse(sb)
+    deps: list[dict] = []
+    seen: set[tuple] = set()
+
+    def add(src: str, dst: str, kind: str):
+        if not dst:
+            return
+        key = (src, dst, kind)
+        if key not in seen:
+            seen.add(key)
+            deps.append({"repo": repo, "path": path, "language": language,
+                         "src": src, "dst": dst, "kind": kind})
+
+    module_name = path.rsplit("/", 1)[-1].rsplit(".", 1)[0] if path else "__module__"
+
+    def walk(node):
+        t = node.type
+        if t in _FUNCTION_TYPES:
+            owner = _sym_name(node, sb)
+            body = node.child_by_field_name("body")
+            if body is not None and owner:
+                _collect_calls(body, owner, sb, add)
+            return  # не спускаемся в вложенные scope как отдельные
+        if t in _CLASS_TYPES:
+            owner = _sym_name(node, sb)
+            if owner:
+                _collect_inherits(node, owner, sb, add)
+            body = node.child_by_field_name("body")
+            if body is not None:
+                for ch in body.children:
+                    walk(ch)  # методы класса — вложенные function_definition
+            return
+        for ch in node.children:
+            walk(ch)
+
+    _collect_imports(tree.root_node, module_name, sb, add)
+    walk(tree.root_node)
+    return deps
+
+
+def store_dependencies(conn, source: str, language: str, repo: str = "", path: str = "",
+                      sync_memgraph: bool = True) -> int:
+    """Извлечь и сохранить зависимости кода (+ рёбра Memgraph).
+
+    Возвращает число сохранённых.
+    """
+    from . import memgraph_store
+    deps = extract_dependencies(source, language, repo, path)
+    n = 0
+    with conn.cursor() as cur:
+        for d in deps:
+            cur.execute(
+                """INSERT INTO pm_code_dependency (repo, path, language, src, dst, kind)
+                   VALUES (%s, %s, %s, %s, %s, %s)
+                   ON CONFLICT (repo, path, src, dst, kind) DO NOTHING""",
+                (d["repo"], d["path"], d["language"], d["src"], d["dst"], d["kind"]),
+            )
+            n += 1
+            if sync_memgraph:
+                try:
+                    memgraph_store.upsert_code_dependency(
+                        d["src"], d["dst"], d["kind"], repo=d["repo"], path=d["path"],
+                    )
+                except Exception:
+                    pass
+    conn.commit()
+    return n
+
+
+def store_file_dependencies(conn, path: str, repo: str = "") -> int:
+    """По расширению файла определить язык и сохранить зависимости."""
+    ext = path.rsplit(".", 1)[-1].lower() if "." in path else ""
+    lang = {".py": "python", ".js": "javascript", ".ts": "typescript", ".rs": "rust"}.get("." + ext)
+    if lang is None:
+        return 0
+    source = open(path, encoding="utf-8").read()
+    return store_dependencies(conn, source, lang, repo, path)