core/causal_graph.py

"""core/causal_graph.py — Graphe causal NetworkX"""
import networkx as nx, hashlib, time, sqlite3, os
from typing import List, Dict, Optional
from pathlib import Path

DATA_DIR = Path(os.environ.get("DATA_DIR", "/tmp/vortex_data"))


class CausalGraph:
    def __init__(self):
        self.graph = nx.MultiDiGraph()
        self._db = None

    def _get_db(self):
        if self._db is None:
            DATA_DIR.mkdir(parents=True, exist_ok=True)
            self._db = sqlite3.connect(str(DATA_DIR / "causal.db"), check_same_thread=False)
            self._db.execute(
                "CREATE TABLE IF NOT EXISTS causal_edges "
                "(id TEXT PRIMARY KEY, src TEXT, tgt TEXT, relation TEXT, weight REAL, ts REAL)"
            )
        return self._db

    def add_node(self, nid: str, kind: str, attrs: Optional[Dict] = None):
        self.graph.add_node(nid, kind=kind, **(attrs or {}))

    def add_edge(self, src: str, tgt: str, relation: str = "causes", weight: float = 1.0):
        eid = hashlib.md5(f"{src}{tgt}{relation}".encode()).hexdigest()[:12]
        for n in (src, tgt):
            if n not in self.graph:
                self.graph.add_node(n, kind="unknown")
        self.graph.add_edge(src, tgt, relation=relation, weight=weight)
        db = self._get_db()
        db.execute(
            "INSERT OR REPLACE INTO causal_edges VALUES (?,?,?,?,?,?)",
            (eid, src, tgt, relation, weight, time.time())
        )
        db.commit()
        return True

    def causal_query(self, src: str, tgt: Optional[str] = None, max_depth: int = 3) -> List[Dict]:
        if src not in self.graph:
            return []
        try:
            if tgt and tgt in self.graph:
                paths = list(nx.all_simple_paths(self.graph, src, tgt, cutoff=max_depth))
                return [{"path": p, "length": len(p)} for p in paths[:10]]
            result = []
            for node in nx.bfs_tree(self.graph, src, depth_limit=max_depth).nodes():
                if node != src:
                    result.append({
                        "node": node,
                        "depth": nx.shortest_path_length(self.graph, src, node)
                    })
            return result[:20]
        except Exception as e:
            return [{"error": str(e)}]

    def stats(self):
        return {"nodes": self.graph.number_of_nodes(), "edges": self.graph.number_of_edges()}