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

pegas_memory/memgraph_store.py

@@ -0,0 +1,136 @@
+"""Memgraph — графовый срез (узлы Atom + CodeSymbol, рёбра RELATES).
+
+Memgraph совместим с Bolt-протоколом Neo4j — используем официальный neo4j драйвер.
+Рёбра единообразны: (любой узел)-[:RELATES {kind, inferred, source}]->(любой узел).
+Это позволяет логической инференции работать и над атомами, и над кодом.
+"""
+from __future__ import annotations
+from neo4j import GraphDatabase
+from . import config
+
+_driver = None
+
+def driver():
+    global _driver
+    if _driver is None:
+        _driver = GraphDatabase.driver(
+            f"bolt://{config.MEMGRAPH_HOST}:{config.MEMGRAPH_PORT}"
+        )
+    return _driver
+
+def upsert_atom(atom_id: int, type: str, claim: str | None, status: str) -> None:
+    with driver().session() as s:
+        s.run(
+            "MERGE (a:Atom {id: $id}) SET a.type = $type, a.claim = $claim, a.status = $status",
+            id=atom_id, type=type, claim=claim or "", status=status,
+        )
+
+def upsert_link(src: int, dst: int, kind: str, *, inferred: bool = False,
+                source: str | None = None, confidence: float = 0.5) -> None:
+    with driver().session() as s:
+        s.run(
+            "MATCH (a:Atom {id: $src}), (b:Atom {id: $dst}) "
+            "MERGE (a)-[r:RELATES {kind: $kind}]->(b) "
+            "SET r.inferred = $inferred, r.source = $source, r.confidence = $confidence",
+            src=src, dst=dst, kind=kind, inferred=inferred, source=source, confidence=confidence,
+        )
+
+def delete_atom(atom_id: int) -> None:
+    with driver().session() as s:
+        s.run("MATCH (a:Atom {id: $id}) DETACH DELETE a", id=atom_id)
+
+def upsert_code_symbol(name: str, *, repo: str = "", path: str = "",
+                       language: str = "", kind: str = "") -> None:
+    """Узел CodeSymbol (дедуп по name+repo+path)."""
+    with driver().session() as s:
+        s.run(
+            "MERGE (c:CodeSymbol {name: $name, repo: $repo, path: $path}) "
+            "SET c.language = $language, c.kind = $kind",
+            name=name, repo=repo, path=path, language=language, kind=kind,
+        )
+
+def upsert_code_dependency(src: str, dst: str, kind: str, *, repo: str = "", path: str = "") -> None:
+    """Ребро зависимости кода: (src)-[:RELATES {kind}]->(dst)."""
+    with driver().session() as s:
+        # dst может быть внешним (builtin) — создаём как CodeSymbol без repo/path
+        s.run(
+            "MERGE (a:CodeSymbol {name: $src, repo: $repo, path: $path}) "
+            "MERGE (b:CodeSymbol {name: $dst}) "
+            "MERGE (a)-[r:RELATES {kind: $kind}]->(b) "
+            "SET r.source = 'extraction'",
+            src=src, dst=dst, kind=kind, repo=repo, path=path,
+        )
+
+def apply_rule_one_step(if_kind: str, then_kind: str) -> int:
+    """Одношаговое правило: (a)-[:RELATES {kind: if}]->(b) ⇒ (a)-[:RELATES {kind: then}]->(b).
+
+    Возвращает число созданных рёбер.
+    """
+    with driver().session() as s:
+        result = s.run(
+            "MATCH (a)-[r:RELATES {kind: $if_kind}]->(b) "
+            "OPTIONAL MATCH (a)-[x:RELATES {kind: $then_kind}]->(b) "
+            "WITH a, b, x WHERE x IS NULL "
+            "CREATE (a)-[r2:RELATES {kind: $then_kind, inferred: true, source: 'inference'}]->(b) "
+            "RETURN count(r2) AS n",
+            if_kind=if_kind, then_kind=then_kind,
+        )
+        return result.single()["n"]
+
+def apply_rule_chain(k1: str, k2: str, then_kind: str) -> int:
+    """Композиция: (a)-[:RELATES {kind: k1}]->(b)-[:RELATES {kind: k2}]->(c)
+    ⇒ (a)-[:RELATES {kind: then}]->(c). Транзитивность. Возвращает число созданных."""
+    with driver().session() as s:
+        result = s.run(
+            "MATCH (a)-[:RELATES {kind: $k1}]->(b)-[:RELATES {kind: $k2}]->(c) "
+            "OPTIONAL MATCH (a)-[x:RELATES {kind: $then_kind}]->(c) "
+            "WITH a, b, c, x WHERE x IS NULL "
+            "CREATE (a)-[r2:RELATES {kind: $then_kind, inferred: true, source: 'inference'}]->(c) "
+            "RETURN count(r2) AS n",
+            k1=k1, k2=k2, then_kind=then_kind,
+        )
+        return result.single()["n"]
+
+def count() -> int:
+    with driver().session() as s:
+        return s.run("MATCH (a:Atom) RETURN count(a) AS n").single()["n"]
+
+def count_edges() -> int:
+    with driver().session() as s:
+        return s.run("MATCH ()-[r:RELATES]->() RETURN count(r) AS n").single()["n"]
+
+def count_code_symbols() -> int:
+    with driver().session() as s:
+        return s.run("MATCH (c:CodeSymbol) RETURN count(c) AS n").single()["n"]
+
+def query_edges(kind: str | None = None, limit: int = 50) -> list[dict]:
+    """Список рёбер (опционально по kind)."""
+    with driver().session() as s:
+        if kind:
+            r = s.run(
+                "MATCH (a)-[r:RELATES {kind: $kind}]->(b) "
+                "RETURN a.name AS src, labels(a) AS al, b.name AS dst, labels(b) AS bl, "
+                "r.kind AS kind, r.inferred AS inferred LIMIT $lim",
+                kind=kind, lim=limit,
+            )
+        else:
+            r = s.run(
+                "MATCH (a)-[r:RELATES]->(b) "
+                "RETURN a.name AS src, labels(a) AS al, b.name AS dst, labels(b) AS bl, "
+                "r.kind AS kind, r.inferred AS inferred LIMIT $lim",
+                lim=limit,
+            )
+        out = []
+        for rec in r:
+            d = {"src": rec["src"], "dst": rec["dst"]}
+            if "al" in rec and rec["al"]:
+                d["src_label"] = rec["al"][0]
+            if "bl" in rec and rec["bl"]:
+                d["dst_label"] = rec["bl"][0]
+            for key in ("kind", "inferred"):
+                try:
+                    d[key] = rec[key]
+                except KeyError:
+                    pass
+            out.append(d)
+        return out