myshell-ai-ComfyUI-ShellAgent-Plugin/custom_routes.py

from io import BytesIO
from pprint import pprint
from aiohttp import web
import os
import requests
import folder_paths
import json
import server
from PIL import Image
import time
import execution
import random
import traceback
import uuid
import asyncio
import logging
from urllib.parse import quote
import threading
import hashlib
import aiohttp
from aiohttp import ClientSession, web
import aiofiles
from typing import Dict, List, Union, Any, Optional
from PIL import Image
import copy
import struct
from aiohttp import web, ClientSession, ClientError, ClientTimeout, ClientResponseError
import atexit
from datetime import datetime
import nodes
import traceback
import re
import keyword

from .dependency_checker import resolve_dependencies


WORKFLOW_ROOT = "shellagent/comfy_workflow"

CustomNodeTypeMap = {
    "ShellAgentPluginInputText": "text",
    "ShellAgentPluginInputInteger": "integer",
    "ShellAgentPluginInputFloat": "number",
    "ShellAgentPluginInputImage": "image",
    "ShellAgentPluginInputVideo": "video",
    "ShellAgentPluginSaveImage": "image",
    "ShellAgentPluginSaveVideoVHS": "video",
}

# Regular expression for a valid Python variable name
variable_name_pattern = r'^[a-zA-Z_][a-zA-Z0-9_]*$'

def is_valid_variable_name(name):
    # Check if it matches the pattern and is not a keyword
    if re.match(variable_name_pattern, name) and not keyword.iskeyword(name):
        return True
    return False

def schema_validator(prompt):
    from nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
    input_names = []
    output_names = []
    schemas = {
        "inputs": {},
        "outputs": {}
    }
    for node_id, node_info in prompt.items():
        node_class_type = node_info.get("class_type")
        if node_class_type is None:
            raise NotImplementedError(f"Missing nodes founded, please first install the missing nodes using ComfyUI Manager")
        node_cls = NODE_CLASS_MAPPINGS[node_class_type]
        if hasattr(node_cls, "RELATIVE_PYTHON_MODULE") and node_cls.RELATIVE_PYTHON_MODULE.startswith("custom_nodes.ComfyUI-ShellAgent-Plugin"):
            schema = {}
            if "input_name" in node_info["inputs"]:
                mode = "inputs"
                input_name = node_info["inputs"]["input_name"]
                if input_name not in input_names:
                    input_names.append(input_name)
                else:
                    raise ValueError(f"Duplicated input_name found in node {NODE_DISPLAY_NAME_MAPPINGS[node_class_type]} with ID={node_id}")
                # handle the schema at the same time
                schema["name"] = input_name
                
            elif "output_name" in node_info["inputs"]:
                mode = "outputs"
                output_name = node_info["inputs"]["output_name"]
                if output_name not in output_names:
                    output_names.append(output_name)
                else:
                    raise ValueError(f"Duplicated output_name found in node {NODE_DISPLAY_NAME_MAPPINGS[node_class_type]} with ID={node_id}")
                schema["name"] = output_name
            else:
                # neither input nor output
                continue
            if hasattr(node_cls, "validate"):
                schema = node_cls.validate(**node_info["inputs"])
                # validate schema
                if not is_valid_variable_name(schema["title"]):
                    raise ValueError(f'`{schema["title"]}` is not a valid variable name!')
            else:
                raise NotImplementedError("the validate is not implemented")
            schemas[mode][node_id] = schema
    return schemas
            
                
@server.PromptServer.instance.routes.get("/shellagent/list_workflow") # data same as queue prompt, plus workflow_name
async def shellagent_list_workflow(request):
    workflow_ids = os.listdir(WORKFLOW_ROOT)
    # append the metadata
    data = []
    for workflow_id in workflow_ids:

        metadata_file = os.path.join(WORKFLOW_ROOT, workflow_id, "metadata.json")
        metadata = json.load(open(metadata_file))
        item = {
            "id": workflow_id,
            "metadata": metadata
        }
        data.append(item)
    return web.json_response(data, status=400)
    
@server.PromptServer.instance.routes.post("/shellagent/get_file") # data same as queue prompt, plus workflow_name
async def shellagent_get_file(request):
    data = await request.json()
    assert data["filename"] in [
        "workflow_api.json",
        "dependencies.json",
        "metadata.json",
        "extra_data.json",
        "schemas.json",
    ]
    
    data = json.load(open(os.path.join(WORKFLOW_ROOT, data["workflow_id"], data["filename"])))
    return web.json_response(data, status=400)
    
@server.PromptServer.instance.routes.post("/shellagent/export") # data same as queue prompt, plus workflow_name
async def shellagent_export(request):
    data = await request.json()
    prompt = data["prompt"]
    custom_dependencies = data.get("custom_dependencies",  {
        "models": {},
        "custom_nodes": {}
    })
    # extra_data = data["extra_data"]
    workflow_id = str(uuid.uuid4())
    
    # metadata.json
    # metadata = {
    #     "name": data["workflow_name"],
    #     "workflow_id": workflow_id,
    #     "create_time": datetime.now().strftime('%Y-%m-%d %H:%M:%S')
    # }
    
    return_dict = {}
    status = 200
    try:
        schemas = schema_validator(prompt)
        # custom_node.json
        dependency_results = resolve_dependencies(prompt, custom_dependencies)
        # save_root = os.path.join(WORKFLOW_ROOT, workflow_id)
        # os.makedirs(save_root, exist_ok=True)
        
        # fname_mapping = {
        #     "workflow_api.json": prompt,
        #     "dependencies.json": dependency_results,
        #     # "metadata.json": metadata,
        #     # "extra_data.json": extra_data,
        #     "schemas.json": schemas,
        # }
        
        # for fname, dict_to_save in fname_mapping.items():
        #     with open(os.path.join(save_root, fname), "w") as f:
        #         json.dump(dict_to_save, f, indent=2)
        warning_message = ""
        if dependency_results.get("black_list_nodes", []):
            warning_message = "The following nodes cannot be deployed to myshell:\n"
            for item in dependency_results["black_list_nodes"]:
                warning_message += f"  {item['name']}: {item['reason']}\n"
                
        if len(schemas["inputs"]) + len(schemas["outputs"]) == 0:
            warning_message += f"The workflow contains neither inputs nor outputs!\n"
        
        return_dict = {
            "success": True,
            "dependencies": dependency_results["dependencies"],
            "warning_message": warning_message,
            "schemas": schemas
        }
    except Exception as e:
        status = 400
        return_dict = {
            "success": False,
            "message_detail": str(traceback.format_exc()),
            "message": str(e),
        }
    return web.json_response(return_dict, status=status)