agent-runtime/agent/engine.py

"""Graph Engine: loads graph definitions, instantiates nodes, executes pipelines."""

import importlib
import logging
from pathlib import Path

from .nodes import NODE_REGISTRY
from .process import ProcessManager

log = logging.getLogger("runtime")

GRAPHS_DIR = Path(__file__).parent / "graphs"


def list_graphs() -> list[dict]:
    """List all available graph definitions."""
    graphs = []
    for f in sorted(GRAPHS_DIR.glob("*.py")):
        if f.name.startswith("_"):
            continue
        mod = _load_graph_module(f.stem)
        if mod:
            graphs.append({
                "name": getattr(mod, "NAME", f.stem),
                "description": getattr(mod, "DESCRIPTION", ""),
                "file": f.name,
            })
    return graphs


def load_graph(name: str) -> dict:
    """Load a graph definition by name. Returns the module's attributes as a dict."""
    # Try matching by NAME attribute first, then by filename
    for f in GRAPHS_DIR.glob("*.py"):
        if f.name.startswith("_"):
            continue
        mod = _load_graph_module(f.stem)
        if mod and getattr(mod, "NAME", "") == name:
            return _graph_from_module(mod)
    # Fallback: match by filename stem
    mod = _load_graph_module(name)
    if mod:
        return _graph_from_module(mod)
    raise ValueError(f"Graph '{name}' not found")


def _load_graph_module(stem: str):
    """Import a graph module by stem name."""
    try:
        return importlib.import_module(f".graphs.{stem}", package="agent")
    except (ImportError, ModuleNotFoundError) as e:
        log.error(f"[engine] failed to load graph '{stem}': {e}")
        return None


def _graph_from_module(mod) -> dict:
    """Extract graph definition from a module."""
    return {
        "name": getattr(mod, "NAME", "unknown"),
        "description": getattr(mod, "DESCRIPTION", ""),
        "nodes": getattr(mod, "NODES", {}),
        "edges": getattr(mod, "EDGES", []),
        "conditions": getattr(mod, "CONDITIONS", {}),
        "audit": getattr(mod, "AUDIT", {}),
        "engine": getattr(mod, "ENGINE", "imperative"),
    }


def instantiate_nodes(graph: dict, send_hud, process_manager: ProcessManager = None) -> dict:
    """Create node instances from a graph definition. Returns {role: node_instance}."""
    nodes = {}
    for role, impl_name in graph["nodes"].items():
        cls = NODE_REGISTRY.get(impl_name)
        if not cls:
            log.error(f"[engine] node class not found: {impl_name}")
            continue
        # Thinker and Expert nodes accept process_manager
        if impl_name.startswith("thinker") or impl_name.endswith("_expert"):
            nodes[role] = cls(send_hud=send_hud, process_manager=process_manager)
        else:
            nodes[role] = cls(send_hud=send_hud)
        log.info(f"[engine] {role} = {impl_name} ({cls.__name__})")
    return nodes


def get_graph_for_cytoscape(graph: dict) -> dict:
    """Convert graph definition to Cytoscape-compatible elements for frontend."""
    elements = {"nodes": [], "edges": []}
    for role in graph["nodes"]:
        elements["nodes"].append({"data": {"id": role, "label": role}})
    for edge in graph["edges"]:
        src = edge["from"]
        targets = edge["to"] if isinstance(edge["to"], list) else [edge["to"]]
        edge_type = edge.get("type", "data")
        for tgt in targets:
            elements["edges"].append({
                "data": {
                    "id": f"e-{src}-{tgt}",
                    "source": src,
                    "target": tgt,
                    "edge_type": edge_type,
                    "condition": edge.get("condition", ""),
                    "carries": edge.get("carries", ""),
                    "method": edge.get("method", ""),
                },
            })
    return elements