connectome.py

"""
Agent Connectome Builder — Flood-Filling Agent Mesh (FFAM)

Applies Google Neural Mapping concepts to multi-agent systems:
- Build a complete map of agent communications (like brain connectomics)
- Track information flow through agent networks (like Flood-Filling Networks)
- Detect bottlenecks, hubs, orphans (like SegCLR cell type discovery)
- Generate synthetic agent graphs for training (like MoGen)

Author: HayulaLab — July 2026
Based on: Google Neural Mapping research (FFN, SegCLR, MoGen, LICONN)
"""
import json, time, os, sys, threading
from pathlib import Path
from collections import defaultdict, deque
from datetime import datetime
import hashlib

try:
    import networkx as nx
except ImportError:
    nx = None
    print("[WARN] networkx not installed — graph analysis disabled")

# ─── Configuration ───────────────────────────────────────
LOG_FILE = Path(os.environ.get("CONNECTOME_LOG", "/tmp/agent-connectome.jsonl"))
SNAPSHOT_DIR = Path(os.environ.get("CONNECTOME_DIR", "/tmp/agent-connectome-snapshots"))
GRAPH_EXPORT = Path(os.environ.get("CONNECTOME_GRAPH", "/tmp/agent-connectome-graph.json"))
FLUSH_INTERVAL = int(os.environ.get("CONNECTOME_FLUSH_MS", "5000"))  # ms

# ─── Event Types (like synapse types) ────────────────────
EVENT_TYPES = {
    "task:dispatch": "excitatory",    # task assigned
    "task:complete": "signal",        # task finished
    "agent:query": "request",         # one agent asks another
    "agent:response": "response",     # reply
    "skill:invoke": "activation",     # skill used
    "memory:read": "read",            # memory access
    "memory:write": "write",          # memory update
    "router:decision": "route",       # routing choice
    "error:timeout": "failure",       # timeout
    "error:refusal": "refusal",       # refusal
}

# ─── Core: Connectome Builder ─────────────────────────────
class AgentConnectome:
    """The complete connectome of a multi-agent system."""
    
    def __init__(self):
        self.agents: dict[str, dict] = {}           # agent_id → metadata
        self.skills: dict[str, dict] = {}            # skill_id → metadata
        self.edges: list[dict] = []                  # communication events
        self.metrics: dict = defaultdict(int)        # aggregate counts
        self.communities: dict = {}                  # detected communities
        self.bottlenecks: list = []                  # detected bottlenecks
        self._lock = threading.Lock()
        self._start_time = time.time()
        
    # ─── Event ingestion (like FFN's voxel classifier) ────
    def ingest(self, event: dict):
        """Record one agent communication event."""
        with self._lock:
            event["_ts"] = time.time()
            event["_idx"] = len(self.edges)
            
            # Register agents
            for field in ["from_agent", "to_agent", "agent"]:
                a = event.get(field)
                if a and a not in self.agents:
                    self.agents[a] = {
                        "id": a,
                        "first_seen": event["_ts"],
                        "events_sent": 0,
                        "events_received": 0,
                        "skills_used": set(),
                        "type": "unknown"
                    }
            
            sender = event.get("from_agent")
            receiver = event.get("to_agent")
            etype = event.get("type", "unknown")
            
            if sender:
                if sender in self.agents:
                    self.agents[sender]["events_sent"] += 1
            if receiver:
                if receiver in self.agents:
                    self.agents[receiver]["events_received"] += 1
            
            skill = event.get("skill")
            if skill:
                if skill not in self.skills:
                    self.skills[skill] = {"id": skill, "invocations": 0, "agents": set()}
                self.skills[skill]["invocations"] += 1
                if sender:
                    self.skills[skill]["agents"].add(sender)
                    if sender in self.agents:
                        self.agents[sender]["skills_used"].add(skill)
            
            self.metrics[f"events:{etype}"] += 1
            self.metrics["total_events"] += 1
            self.edges.append(event)
    
    # ─── Build graph (like connectome reconstruction) ─────
    def build_graph(self) -> dict:
        """Build full agent connectome."""
        G = nx.DiGraph()
        
        for aid, adata in self.agents.items():
            G.add_node(aid, **adata)
        
        edge_weights = defaultdict(int)
        for e in self.edges:
            u, v = e.get("from_agent"), e.get("to_agent")
            if u and v:
                edge_weights[(u, v)] += 1
                edge_weights[(v, u)] += 0  # track reverse
        
        for (u, v), w in edge_weights.items():
            if w > 0:
                etype = "bidirectional" if edge_weights.get((v, u), 0) > 0 else "unidirectional"
                G.add_edge(u, v, weight=w, type=etype)
        
        return {
            "nodes": len(G.nodes),
            "edges": len(G.edges),
            "density": nx.density(G) if len(G) > 1 else 0,
            "is_connected": nx.is_weakly_connected(G) if len(G) > 1 else False,
            "diameter": nx.diameter(G.to_undirected()) if len(G) > 1 and nx.is_connected(G.to_undirected()) else -1,
            "avg_path_length": nx.average_shortest_path_length(G.to_undirected()) if len(G) > 1 and nx.is_connected(G.to_undirected()) else -1,
        }
    
    # ─── Hub detection (like SegCLR cell type discovery) ──
    def find_hubs(self, min_connections: int = 3) -> list[dict]:
        """Find hub agents (most connected) — like hub neurons."""
        G = nx.DiGraph()
        for aid in self.agents:
            G.add_node(aid)
        for e in self.edges:
            u, v = e.get("from_agent"), e.get("to_agent")
            if u and v:
                G.add_edge(u, v)
        
        hubs = []
        for node in G.nodes():
            degree = G.degree(node)
            in_deg = G.in_degree(node)
            out_deg = G.out_degree(node)
            if degree >= min_connections:
                hubs.append({
                    "agent": node,
                    "degree": degree,
                    "in_degree": in_deg,
                    "out_degree": out_deg,
                    "betweenness": nx.betweenness_centrality(G).get(node, 0),
                    "type": "router" if out_deg > in_deg * 2 else
                            "aggregator" if in_deg > out_deg * 2 else
                            "peer"
                })
        hubs.sort(key=lambda x: x["degree"], reverse=True)
        return hubs
    
    # ─── Bottleneck detection ──────────────────────────────
    def find_bottlenecks(self, threshold: float = 0.3) -> list[dict]:
        """Find bottlenecks — agents that are single points of failure."""
        G = nx.DiGraph()
        for aid in self.agents:
            G.add_node(aid)
        for e in self.edges:
            u, v = e.get("from_agent"), e.get("to_agent")
            if u and v:
                G.add_edge(u, v)
        
        if len(G) < 3:
            return []
        
        try:
            bc = nx.betweenness_centrality(G)
            avg_bc = sum(bc.values()) / len(bc) if bc else 0
            bottlenecks = []
            for node, score in bc.items():
                if score > avg_bc * (1 + threshold):
                    bottlenecks.append({
                        "agent": node,
                        "betweenness": score,
                        "severity": "critical" if score > avg_bc * 3 else "high" if score > avg_bc * 2 else "moderate",
                        "recommendation": "Add redundant agent" if score > avg_bc * 3 else
                                        "Consider load balancing" if score > avg_bc * 2 else
                                        "Monitor"
                    })
            return sorted(bottlenecks, key=lambda x: x["betweenness"], reverse=True)
        except:
            return []
    
    # ─── Critical path analysis (like neural pathway tracing) ──
    def critical_paths(self, top_k: int = 5) -> list[dict]:
        """Find the most common agent chains (critical paths)."""
        paths = defaultdict(int)
        
        # Build agent sequences from events
        sequences = []
        current_seq = []
        for e in self.edges:
            sender = e.get("from_agent")
            receiver = e.get("to_agent")
            if sender:
                if not current_seq or current_seq[-1] != sender:
                    current_seq.append(sender)
            if receiver:
                current_seq.append(receiver)
        
        # Find common subsequences
        for i in range(len(current_seq)):
            for j in range(i+2, min(i+8, len(current_seq))):
                seq = tuple(current_seq[i:j])
                paths[seq] += 1
        
        top = sorted(paths.items(), key=lambda x: x[1], reverse=True)[:top_k]
        return [{"path": list(p), "frequency": f} for p, f in top]
    
    # ─── Synthetic graph generation (MoGen-inspired) ──────
    def generate_synthetic(self, num_agents: int = 10, density: float = 0.3) -> list[dict]:
        """Generate synthetic agent graphs for training — like MoGen's synthetic neurons."""
        if not nx:
            return []
        
        G = nx.gnp_random_graph(num_agents, density, directed=True)
        agents = []
        for i in range(num_agents):
            agent_type = nx.random.choice(["router", "worker", "verifier", "memory", "observer"], 
                                         p=[0.15, 0.5, 0.1, 0.15, 0.1])
            agents.append({
                "id": f"synth-agent-{i:03d}",
                "type": agent_type,
                "connections": list(G.neighbors(i)),
                "degree": G.degree(i),
            })
        return agents
    
    # ─── Snapshot (like Neuroglancer scene capture) ───────
    def snapshot(self) -> dict:
        """Take a complete snapshot of the connectome."""
        return {
            "timestamp": datetime.now().isoformat(),
            "uptime_seconds": time.time() - self._start_time,
            "stats": {
                "agents": len(self.agents),
                "skills": len(self.skills),
                "events": len(self.edges),
                "metrics": dict(self.metrics),
            },
            "graph": self.build_graph(),
            "hubs": self.find_hubs(),
            "bottlenecks": self.find_bottlenecks(),
            "critical_paths": self.critical_paths(),
            "agent_list": list(self.agents.keys()),
            "skill_list": list(self.skills.keys()),
        }
    
    def save(self, path: str = None):
        """Save snapshot to JSON."""
        path = path or str(SNAPSHOT_DIR / f"connectome-{int(time.time())}.json")
        snapshot = self.snapshot()
        Path(path).parent.mkdir(parents=True, exist_ok=True)
        with open(path, "w") as f:
            json.dump(snapshot, f, indent=2, ensure_ascii=False)
        return path


# ─── Integration: Hook into existing Hayula infrastructure ──
class ConnectomeIntegrator:
    """Hooks the connectome into DragonMesh, EventBus, Observability."""
    
    def __init__(self):
        self.connectome = AgentConnectome()
        self._running = False
        self._thread = None
    
    def hook_eventbus(self, eventbus):
        """Wrap EventBus.publish to record all events."""
        original_publish = eventbus.publish
        
        def traced_publish(event):
            self.connectome.ingest(event)
            return original_publish(event)
        
        eventbus.publish = traced_publish
        return eventbus
    
    def hook_dragonmesh(self, mesh):
        """Wrap DragonMesh route to trace routing decisions."""
        if hasattr(mesh, 'route'):
            original_route = mesh.route
            def traced_route(task):
                result = original_route(task)
                self.connectome.ingest({
                    "type": "router:decision",
                    "from_agent": "dragon_mesh",
                    "to_agent": result.get("agent", "unknown"),
                    "task": str(task)[:100]
                })
                return result
            mesh.route = traced_route
        return mesh
    
    def hook_a2a(self, bridge):
        """Wrap A2A bridge to trace agent-to-agent communication."""
        if hasattr(bridge, 'send'):
            original_send = bridge.send
            def traced_send(agent, message):
                self.connectome.ingest({
                    "type": "agent:query",
                    "from_agent": "bridge",
                    "to_agent": agent,
                    "message": str(message)[:200]
                })
                result = original_send(agent, message)
                self.connectome.ingest({
                    "type": "agent:response",
                    "from_agent": agent,
                    "to_agent": "bridge",
                    "result": str(result)[:200]
                })
                return result
            bridge.send = traced_send
        return bridge


# ─── CLI ──────────────────────────────────────────────────
if __name__ == "__main__":
    import argparse
    p = argparse.ArgumentParser(description="Agent Connectome — Flood-Filling Agent Mesh")
    sp = p.add_subparsers(dest="cmd")
    
    # Demo: simulate agent traffic
    sp.add_parser("demo", help="Run demo with simulated agent traffic")
    
    # Analyze existing log
    analyze = sp.add_parser("analyze", help="Analyze agent connectome from log")
    analyze.add_argument("--log", default=str(LOG_FILE))
    
    # Generate synthetic graph
    synth = sp.add_parser("synth", help="Generate synthetic agent graph")
    synth.add_argument("-n", type=int, default=10, help="Number of synthetic agents")
    synth.add_argument("-d", type=float, default=0.3, help="Graph density")
    
    # Snapshot
    sp.add_parser("snapshot", help="Take connectome snapshot")
    
    args = p.parse_args()
    
    if args.cmd == "demo":
        connectome = AgentConnectome()
        agents = ["rushd", "wafa", "awf", "dragon", "hermes", "musa", "zeus", "haytham"]
        skills = ["code_review", "text_gen", "trade_signal", "memory_search", "task_route"]
        
        print(f"[FFAM] Starting demo with {len(agents)} agents, {len(skills)} skills")
        
        for i in range(100):
            import random
            sender = random.choice(agents)
            receiver = random.choice([a for a in agents if a != sender])
            event = {
                "type": random.choice(list(EVENT_TYPES.keys())),
                "from_agent": sender,
                "to_agent": receiver,
                "skill": random.choice(skills) if random.random() > 0.5 else None,
                "task_id": f"task-{i:04d}",
            }
            connectome.ingest(event)
            time.sleep(0.01)
        
        snap = connectome.snapshot()
        print(json.dumps(snap["stats"], indent=2))
        print(f"\n🔍 Hubs:")
        for h in snap["hubs"][:5]:
            print(f"  {h['agent']:12} degree={h['degree']:3d} type={h['type']}")
        print(f"\n⚠️  Bottlenecks:")
        for b in snap["bottlenecks"][:3]:
            print(f"  {b['agent']:12} severity={b['severity']:10} → {b['recommendation']}")
        print(f"\n🛤️  Critical paths:")
        for cp in snap["critical_paths"]:
            print(f"  {' → '.join(cp['path'])}  (×{cp['frequency']})")
        
        connectome.save()
        print(f"\n✅ Snapshot saved to {SNAPSHOT_DIR}")
        
    elif args.cmd == "synth":
        connectome = AgentConnectome()
        g = connectome.generate_synthetic(args.n, args.d)
        print(json.dumps(g, indent=2))
        
    elif args.cmd == "snapshot":
        connectome = AgentConnectome()
        if LOG_FILE.exists():
            with open(LOG_FILE) as f:
                for line in f:
                    connectome.ingest(json.loads(line.strip()))
        path = connectome.save()
        print(f"Snap: {path}")
        print(json.dumps(connectome.snapshot()["stats"], indent=2))