73 lines
2.7 KiB
Python
73 lines
2.7 KiB
Python
|
from fastapi import APIRouter, HTTPException, Body
|
||
|
|
from pytrace.core.registry import NodeRegistry, get_type_string
|
||
|
|
import importlib.util
|
||
|
|
import sys
|
||
|
|
from pathlib import Path
|
||
|
|
import os
|
||
|
|
|
||
|
|
router = APIRouter()
|
||
|
|
|
||
|
|
# Create a directory to store custom nodes uploaded by clients
|
||
|
|
CUSTOM_NODES_DIR = Path("custom_nodes")
|
||
|
|
CUSTOM_NODES_DIR.mkdir(exist_ok=True)
|
||
|
|
if str(CUSTOM_NODES_DIR.resolve()) not in sys.path:
|
||
|
|
sys.path.insert(0, str(CUSTOM_NODES_DIR.resolve()))
|
||
|
|
|
||
|
|
|
||
|
|
@router.post("/register")
|
||
|
|
async def register_custom_node(code: str = Body(..., embed=True)):
|
||
|
|
"""
|
||
|
|
Receives Python code for a custom node, saves it, and loads it.
|
||
|
|
"""
|
||
|
|
# This is a simplified and dangerous implementation for demonstration.
|
||
|
|
# A real implementation would need heavy security validation (AST analysis, sandboxing).
|
||
|
|
|
||
|
|
try:
|
||
|
|
# For simplicity, we'll hash the code to create a filename.
|
||
|
|
import hashlib
|
||
|
|
code_hash = hashlib.sha256(code.encode()).hexdigest()
|
||
|
|
file_path = CUSTOM_NODES_DIR / f"node_{code_hash}.py"
|
||
|
|
|
||
|
|
if file_path.exists():
|
||
|
|
# Node is already registered, we can assume.
|
||
|
|
return {"message": "Node already exists and is registered."}
|
||
|
|
|
||
|
|
file_path.write_text(code, encoding="utf-8")
|
||
|
|
|
||
|
|
# Dynamically load the module.
|
||
|
|
# The node's class should have the `@register_node` decorator,
|
||
|
|
# so importing it will automatically register it.
|
||
|
|
module_name = file_path.stem
|
||
|
|
spec = importlib.util.spec_from_file_location(module_name, file_path)
|
||
|
|
module = importlib.util.module_from_spec(spec)
|
||
|
|
spec.loader.exec_module(module)
|
||
|
|
|
||
|
|
# Find which node was just registered by looking for a new class
|
||
|
|
# from the loaded module.
|
||
|
|
newly_registered_type = None
|
||
|
|
for node_type, node_class in NodeRegistry._implementations.items():
|
||
|
|
if node_class.__module__ == module_name:
|
||
|
|
newly_registered_type = node_type
|
||
|
|
break
|
||
|
|
|
||
|
|
if not newly_registered_type:
|
||
|
|
# Cleanup the created file if registration fails
|
||
|
|
os.remove(file_path)
|
||
|
|
raise HTTPException(status_code=400, detail="Could not find a class with @register_node in the provided code.")
|
||
|
|
|
||
|
|
return {
|
||
|
|
"message": "Node registered successfully.",
|
||
|
|
"node_type": newly_registered_type,
|
||
|
|
"spec": NodeRegistry.get_spec(newly_registered_type).dict()
|
||
|
|
}
|
||
|
|
|
||
|
|
except Exception as e:
|
||
|
|
import traceback
|
||
|
|
raise HTTPException(
|
||
|
|
status_code=500,
|
||
|
|
detail={
|
||
|
|
"error": str(e),
|
||
|
|
"traceback": traceback.format_exc()
|
||
|
|
}
|
||
|
|
)
|