137 lines
5.9 KiB
Python
137 lines
5.9 KiB
Python
"""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
|