diff --git a/__init__.py b/__init__.py index 121af4d..39b03a4 100644 --- a/__init__.py +++ b/__init__.py @@ -1,19 +1,41 @@ -from .src.core.llm_loader import NODE_CLASS_MAPPINGS as LLM_LOADER_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as LLM_LOADER_DISPLAY_MAPPINGS -from .src.core.openai_compatible import NODE_CLASS_MAPPINGS as OPENAI_COMPATIBLE_LOADER_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as OPENAI_COMPATIBLE_LOADER_DISPLAY_MAPPINGS -from .src.processors.image_processor import NODE_CLASS_MAPPINGS as IMAGE_PREPROCESSOR_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as IMAGE_PREPROCESSOR_DISPLAY_MAPPINGS +""" +ComfyUI-LLMs-Toolkit: Lightweight LLM API integration for ComfyUI. -# 合并所有节点映射 -NODE_CLASS_MAPPINGS = { - **LLM_LOADER_MAPPINGS, - **OPENAI_COMPATIBLE_LOADER_MAPPINGS, - **IMAGE_PREPROCESSOR_MAPPINGS, -} +Auto-registration mechanism for all nodes. +""" +from pathlib import Path +import importlib.util +import sys -NODE_DISPLAY_NAME_MAPPINGS = { - **LLM_LOADER_DISPLAY_MAPPINGS, - **OPENAI_COMPATIBLE_LOADER_DISPLAY_MAPPINGS, - **IMAGE_PREPROCESSOR_DISPLAY_MAPPINGS, -} +# Initialize mappings +NODE_CLASS_MAPPINGS = {} +NODE_DISPLAY_NAME_MAPPINGS = {} + +# Auto-load all nodes from nodes/ directory +nodes_dir = Path(__file__).parent / "nodes" + +if nodes_dir.exists(): + for py_file in sorted(nodes_dir.glob("*.py")): + if py_file.stem.startswith("_"): + continue + + try: + # Load module + module_name = f"ComfyUI-LLMs-Toolkit.nodes.{py_file.stem}" + spec = importlib.util.spec_from_file_location(module_name, py_file) + module = importlib.util.module_from_spec(spec) + sys.modules[module_name] = module + spec.loader.exec_module(module) + + # Register node mappings + if hasattr(module, "NODE_CLASS_MAPPINGS"): + NODE_CLASS_MAPPINGS.update(module.NODE_CLASS_MAPPINGS) + + if hasattr(module, "NODE_DISPLAY_NAME_MAPPINGS"): + NODE_DISPLAY_NAME_MAPPINGS.update(module.NODE_DISPLAY_NAME_MAPPINGS) + + except Exception as e: + print(f"[LLMs_Toolkit] Failed to load {py_file.stem}: {e}") WEB_DIRECTORY = "./web" __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"] diff --git a/src/processors/image_processor.py b/nodes/image_processor.py similarity index 100% rename from src/processors/image_processor.py rename to nodes/image_processor.py diff --git a/nodes/json_extractor.py b/nodes/json_extractor.py new file mode 100644 index 0000000..3992639 --- /dev/null +++ b/nodes/json_extractor.py @@ -0,0 +1,78 @@ +import json +from typing import Tuple + + +class JSONExtractor: + """ + Extract value from JSON string by key. + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "json_string": ("STRING", { + "multiline": True, + "default": "{}" + }), + "key": ("STRING", { + "default": "", + "label": "Key to extract" + }) + } + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("value",) + FUNCTION = "extract" + CATEGORY = "🚦ComfyUI_LLMs_Toolkit/JSON" + + def extract(self, json_string: str, key: str) -> Tuple[str]: + """ + Extract value from JSON string by key. + + Args: + json_string: JSON string to parse + key: Key to extract (supports nested keys like 'user.name') + + Returns: + Extracted value as string + """ + try: + # Parse JSON string + data = json.loads(json_string) + + # Handle nested keys (e.g., "user.profile.name") + keys = key.split('.') + value = data + + for k in keys: + if isinstance(value, dict): + value = value.get(k) + if value is None: + return (f"Key '{k}' not found",) + else: + return (f"Cannot access key '{k}' on non-dict value",) + + # Convert value to string + if isinstance(value, (dict, list)): + result = json.dumps(value, ensure_ascii=False) + else: + result = str(value) + + print(f"[LLMs_Toolkit] extracted {key}={result[:50]}...") + return (result,) + + except json.JSONDecodeError as e: + error_msg = f"Invalid JSON: {str(e)}" + print(f"[LLMs_Toolkit] ✗ {error_msg}") + return (error_msg,) + except Exception as e: + error_msg = f"Extraction error: {str(e)}" + print(f"[LLMs_Toolkit] ✗ {error_msg}") + return (error_msg,) + + +# Register the node +NODE_CLASS_MAPPINGS = {"JSONExtractor": JSONExtractor} +NODE_DISPLAY_NAME_MAPPINGS = {"JSONExtractor": "JSON Extractor"} diff --git a/src/core/llm_loader.py b/nodes/llm_loader.py similarity index 100% rename from src/core/llm_loader.py rename to nodes/llm_loader.py diff --git a/src/core/openai_compatible.py b/nodes/openai_compatible.py similarity index 100% rename from src/core/openai_compatible.py rename to nodes/openai_compatible.py diff --git a/src/__init__.py b/src/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/src/core/__init__.py b/src/core/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/src/nodes/__init__.py b/src/nodes/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/src/processors/__init__.py b/src/processors/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/src/utils/__init__.py b/src/utils/__init__.py deleted file mode 100644 index 3a81e55..0000000 --- a/src/utils/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ -from . import file_uploader -__all__ = ["file_uploader"]