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

pegas_memory/code_symbols.py

@@ -0,0 +1,108 @@
+"""Извлечение символов кода через tree-sitter (Фаза 3).
+
+Грамматики ставятся отдельными пакетами: tree-sitter-python, tree-sitter-javascript,
+tree-sitter-rust и т.д. Каждый пакет экспортирует language().
+"""
+from __future__ import annotations
+import hashlib
+
+# Язык -> (имя модуля, признак function/class узлов)
+_LANGS = {
+    "python": ("tree_sitter_python", {"function_definition", "class_definition"}),
+    "javascript": ("tree_sitter_javascript", {"function_declaration", "class_declaration", "method_definition"}),
+    "typescript": ("tree_sitter_javascript", {"function_declaration", "class_declaration", "method_definition"}),
+    "rust": ("tree_sitter_rust", {"function_item", "struct_item", "impl_item", "trait_item", "enum_item"}),
+}
+
+
+def _parser(language: str):
+    """Вернуть (Parser, kind_nodes) для языка."""
+    from tree_sitter import Language, Parser
+    spec = _LANGS.get(language)
+    if spec is None:
+        raise ValueError(f"unsupported language: {language}")
+    mod_name, kind_nodes = spec
+    mod = __import__(mod_name)
+    lang = Language(mod.language())
+    return Parser(lang), kind_nodes
+
+
+def extract_symbols(source: str, language: str, repo: str = "", path: str = "") -> list[dict]:
+    """Вернуть символы (function/class/method/type/const) из исходника.
+
+    tree-sitter индексирует байтами UTF-8, поэтому слайсим из bytes (src),
+    а не из str (source) — иначе кириллица в комментариях сбивает индексы.
+    """
+    parser, kind_nodes = _parser(language)
+    src = source.encode("utf-8")
+    tree = parser.parse(src)
+    symbols: list[dict] = []
+
+    def walk(node):
+        if node.type in kind_nodes:
+            name_node = node.child_by_field_name("name")
+            name = src[name_node.start_byte:name_node.end_byte].decode("utf-8", "replace") if name_node else "?"
+            # сигнатура: от начала узла до первого тела/блока (первые 200 байт)
+            sig = src[node.start_byte:min(node.end_byte, node.start_byte + 200)].decode("utf-8", "replace")
+            if "function" in node.type or node.type in ("function_item", "function_declaration", "method_definition"):
+                kind = "function"
+            elif "class" in node.type or node.type in ("struct_item", "impl_item", "trait_item"):
+                kind = "class"
+            elif "enum" in node.type:
+                kind = "type"
+            else:
+                kind = "type"
+            symbols.append({
+                "repo": repo, "path": path, "language": language,
+                "symbol_name": name, "kind": kind,
+                "signature": sig,
+                "body_hash": hashlib.sha256(src[node.start_byte:node.end_byte]).hexdigest(),
+                "line_start": node.start_point[0] + 1, "line_end": node.end_point[0] + 1,
+            })
+        for child in node.children:
+            walk(child)
+
+    walk(tree.root_node)
+    return symbols
+
+
+def store_symbols(conn, source: str, language: str, repo: str = "", path: str = "",
+                 sync_memgraph: bool = True) -> int:
+    """Извлечь и сохранить символы кода в pm_code_symbol (+ узлы Memgraph).
+
+    Возвращает число сохранённых.
+    """
+    from . import memgraph_store
+    syms = extract_symbols(source, language, repo, path)
+    n = 0
+    with conn.cursor() as cur:
+        for s in syms:
+            cur.execute(
+                """INSERT INTO pm_code_symbol (repo, path, language, symbol_name, kind,
+                                               signature, body_hash, line_start, line_end)
+                   VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s)
+                   ON CONFLICT DO NOTHING""",
+                (s["repo"], s["path"], s["language"], s["symbol_name"], s["kind"],
+                 s["signature"], s["body_hash"], s["line_start"], s["line_end"]),
+            )
+            n += 1
+            if sync_memgraph:
+                try:
+                    memgraph_store.upsert_code_symbol(
+                        s["symbol_name"], repo=s["repo"], path=s["path"],
+                        language=s["language"], kind=s["kind"],
+                    )
+                except Exception:
+                    pass
+    conn.commit()
+    return n
+
+
+def store_file(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_symbols(conn, source, lang, repo, path)