unreal-engine-mcp-system-plugin/ue_side/ai.py

"""AI ops — BehaviorTree, Blackboard, EQS, AIController.

Ops:
  create_behavior_tree    {path, blackboard?}
  create_blackboard       {path, parent?}
  create_eqs_query        {path}
  create_aicontroller_bp  {path, parent_class?:'AIController'}
  add_blackboard_key      {path, name, type:'Bool|Int|Float|String|Name|Vector|Rotator|Object|Class|Enum', base_class?}
  remove_blackboard_key   {path, name}
  list_blackboard_keys    {path}
  set_bt_blackboard       {bt_path, blackboard_path}
"""
from __future__ import annotations

import json
import os
import traceback

try:
    import unreal  # type: ignore
except ImportError:
    unreal = None


def _at(): return unreal.AssetToolsHelpers.get_asset_tools()
def _load(p): return unreal.EditorAssetLibrary.load_asset(p)
def _split(p):
    folder, name = p.rsplit("/", 1)
    return folder, name


KEY_TYPES = {
    "Bool":    "BlackboardKeyType_Bool",
    "Int":     "BlackboardKeyType_Int",
    "Float":   "BlackboardKeyType_Float",
    "String":  "BlackboardKeyType_String",
    "Name":    "BlackboardKeyType_Name",
    "Vector":  "BlackboardKeyType_Vector",
    "Rotator": "BlackboardKeyType_Rotator",
    "Object":  "BlackboardKeyType_Object",
    "Class":   "BlackboardKeyType_Class",
    "Enum":    "BlackboardKeyType_Enum",
}


def create_behavior_tree(args):
    folder, name = _split(args["path"])
    factory_cls = getattr(unreal, "BehaviorTreeFactory", None)
    if factory_cls is None:
        return {"ok": False, "error": "BehaviorTreeFactory missing — enable AIModule"}
    asset = _at().create_asset(name, folder, unreal.BehaviorTree, factory_cls())
    if asset is None:
        return {"ok": False, "error": "create returned None"}
    if args.get("blackboard"):
        bb = _load(args["blackboard"])
        if bb is not None:
            asset.set_editor_property("blackboard_asset", bb)
    unreal.EditorAssetLibrary.save_asset(asset.get_path_name())
    return {"ok": True, "path": asset.get_path_name().split(".")[0]}


def create_blackboard(args):
    folder, name = _split(args["path"])
    factory_cls = getattr(unreal, "BlackboardDataFactory", None)
    if factory_cls is None:
        return {"ok": False, "error": "BlackboardDataFactory missing"}
    asset = _at().create_asset(name, folder, unreal.BlackboardData, factory_cls())
    if asset is None:
        return {"ok": False, "error": "create returned None"}
    if args.get("parent"):
        parent = _load(args["parent"])
        if parent is not None:
            try:
                asset.set_editor_property("parent", parent)
            except Exception:
                pass
    unreal.EditorAssetLibrary.save_asset(asset.get_path_name())
    return {"ok": True, "path": asset.get_path_name().split(".")[0]}


def create_eqs_query(args):
    folder, name = _split(args["path"])
    factory_cls = getattr(unreal, "EnvQueryFactory", None)
    asset_cls = getattr(unreal, "EnvQuery", None)
    if factory_cls is None or asset_cls is None:
        return {"ok": False, "error": "EnvQuery factory missing — enable EQS in project settings"}
    asset = _at().create_asset(name, folder, asset_cls, factory_cls())
    if asset is None:
        return {"ok": False, "error": "create returned None"}
    unreal.EditorAssetLibrary.save_asset(asset.get_path_name())
    return {"ok": True, "path": asset.get_path_name().split(".")[0]}


def create_aicontroller_bp(args):
    folder, name = _split(args["path"])
    parent_name = args.get("parent_class") or "AIController"
    parent = unreal.load_class(None, "/Script/AIModule." + parent_name)
    if parent is None:
        parent = unreal.load_class(None, parent_name)
    if parent is None:
        return {"ok": False, "error": f"parent {parent_name!r} not found"}
    factory = unreal.BlueprintFactory()
    factory.set_editor_property("parent_class", parent)
    asset = _at().create_asset(name, folder, unreal.Blueprint, factory)
    if asset is None:
        return {"ok": False, "error": "create returned None"}
    unreal.EditorAssetLibrary.save_asset(asset.get_path_name())
    return {"ok": True, "path": asset.get_path_name().split(".")[0]}


def add_blackboard_key(args):
    bb = _load(args["path"])
    if bb is None:
        return {"ok": False, "error": "blackboard not found"}
    type_name = KEY_TYPES.get(args["type"])
    if not type_name:
        return {"ok": False, "error": f"unknown key type {args['type']!r}; valid: {list(KEY_TYPES)}"}
    key_type_cls = getattr(unreal, type_name, None)
    if key_type_cls is None:
        return {"ok": False, "error": f"{type_name} not exposed"}
    entry_cls = getattr(unreal, "BlackboardEntry", None)
    if entry_cls is None:
        return {"ok": False, "error": "BlackboardEntry not exposed; edit via UI"}
    entry = entry_cls()
    entry.set_editor_property("entry_name", args["name"])
    kt = key_type_cls()
    if args.get("base_class") and hasattr(kt, "base_class"):
        cls = unreal.load_class(None, args["base_class"])
        if cls is not None:
            kt.set_editor_property("base_class", cls)
    entry.set_editor_property("key_type", kt)
    keys = list(bb.get_editor_property("keys") or [])
    keys.append(entry)
    bb.set_editor_property("keys", keys)
    unreal.EditorAssetLibrary.save_asset(bb.get_path_name())
    return {"ok": True, "name": args["name"]}


def remove_blackboard_key(args):
    bb = _load(args["path"])
    if bb is None:
        return {"ok": False, "error": "blackboard not found"}
    target = args["name"]
    keys = [k for k in (bb.get_editor_property("keys") or [])
            if str(k.get_editor_property("entry_name")) != target]
    bb.set_editor_property("keys", keys)
    unreal.EditorAssetLibrary.save_asset(bb.get_path_name())
    return {"ok": True, "removed": target}


def list_blackboard_keys(args):
    bb = _load(args["path"])
    if bb is None:
        return {"ok": False, "error": "blackboard not found"}
    out = []
    for k in (bb.get_editor_property("keys") or []):
        try:
            kt = k.get_editor_property("key_type")
            out.append({"name": str(k.get_editor_property("entry_name")),
                        "type": kt.get_class().get_name() if kt else None})
        except Exception:
            pass
    return {"ok": True, "keys": out}


def set_bt_blackboard(args):
    bt = _load(args["bt_path"])
    bb = _load(args["blackboard_path"])
    if bt is None or bb is None:
        return {"ok": False, "error": "bt or blackboard not found"}
    bt.set_editor_property("blackboard_asset", bb)
    unreal.EditorAssetLibrary.save_asset(bt.get_path_name())
    return {"ok": True}


OPS = {
    "create_behavior_tree":   create_behavior_tree,
    "create_blackboard":      create_blackboard,
    "create_eqs_query":       create_eqs_query,
    "create_aicontroller_bp": create_aicontroller_bp,
    "add_blackboard_key":     add_blackboard_key,
    "remove_blackboard_key":  remove_blackboard_key,
    "list_blackboard_keys":   list_blackboard_keys,
    "set_bt_blackboard":      set_bt_blackboard,
}


def entrypoint(payload_json):
    try:
        p = json.loads(payload_json)
    except Exception as e:
        return _tee({"ok": False, "error": f"bad json: {e}"}, {})
    op = p.get("op"); fn = OPS.get(op)
    if not fn:
        return _tee({"ok": False, "error": f"unknown op {op!r}", "available": list(OPS)}, p)
    try:
        return _tee(fn(p), p)
    except Exception as e:
        return _tee({"ok": False, "error": str(e), "traceback": traceback.format_exc()}, p)


def _tee(result, p):
    txt = json.dumps(result)
    try:
        if unreal is not None:
            sd = unreal.Paths.project_saved_dir() + "MCP/"
            os.makedirs(sd, exist_ok=True)
            rid = p.get("__rid") or "last"
            with open(sd + "last.json", "w", encoding="utf-8") as f: f.write(txt)
            if rid != "last":
                with open(sd + f"{rid}.json", "w", encoding="utf-8") as f: f.write(txt)
            unreal.log("[MCP-AI] " + txt[:1500])
    except Exception:
        pass
    return txt