From 35d908e4d9485fe47779119905abefcd99f955a5 Mon Sep 17 00:00:00 2001 From: CentralFloridaAttorney <71934273+CentralFloridaAttorney@users.noreply.github.com> Date: Fri, 3 Apr 2026 13:20:49 -0400 Subject: [PATCH] Save Current State: zmongo nodes refactor --- api/zmongo_record_editor_api.py | 224 ++ api/zmongo_tabular_record_view_api.py | 93 + .../zmongo_chat_agent_starter.jpg | Bin .../zmongo_db_browser.jpg | Bin .../zmongo_loop_sd15_demo.jpg | Bin .../zmongo_record_loop_demoa - Copy.jpg | Bin .../zmongo_record_loop_demoa.jpg | Bin .../zmongo_system_demoa.jpg | Bin .../zmongo_system_demoa.json | 0 nodes/model_loader/__init__.py | 17 + .../zmongo_builtin_loader_adapter_node.py | 437 +++ ...zmongo_model_compatibility_display_node.py | 116 + .../zmongo_model_introspector_node.py | 395 +++ .../zmongo_universal_adapter_node.py | 222 ++ .../zmongo_universal_model_loader_node.py | 378 +++ .../zmongo_universal_model_selector_node.py | 133 + nodes/zmongo_field_selector_node.py | 213 ++ nodes/zmongo_flattened_field_selector_node.py | 213 ++ nodes/zmongo_nodes.py | 2788 +++++++++++++++++ nodes/zmongo_record_editor_node.py | 259 ++ nodes/zmongo_record_splitter.py | 55 + nodes/zmongo_tabular_record_view_node.py | 367 +++ web/js/app.js | 169 + web/js/preset_panel.js | 336 ++ web/js/zmongo_field_selector.js | 148 + web/js/zmongo_field_selector_dropdown.js | 145 + web/js/zmongo_record_editor.js | 332 ++ web/js/zmongo_tabular_record_view.js | 508 +++ web/js/zmongo_universal_adapter.js | 88 + 29 files changed, 7636 insertions(+) create mode 100644 api/zmongo_record_editor_api.py create mode 100644 api/zmongo_tabular_record_view_api.py rename example_workflows/zmongo_db_browser.png => examples/zmongo_chat_agent_starter.jpg (100%) rename example_workflows/zmongo_loop_sd15_demo.png => examples/zmongo_db_browser.jpg (100%) rename example_workflows/zmongo_record_loop_demoa - Copy.jpg => examples/zmongo_loop_sd15_demo.jpg (100%) rename example_workflows/zmongo_record_loop_demoa.png => examples/zmongo_record_loop_demoa - Copy.jpg (100%) rename example_workflows/zmongo_record_loop_demob.jpg => examples/zmongo_record_loop_demoa.jpg (100%) rename example_workflows/zmongo_system_demoa.png => examples/zmongo_system_demoa.jpg (100%) rename {example_workflows => examples}/zmongo_system_demoa.json (100%) create mode 100644 nodes/model_loader/__init__.py create mode 100644 nodes/model_loader/zmongo_builtin_loader_adapter_node.py create mode 100644 nodes/model_loader/zmongo_model_compatibility_display_node.py create mode 100644 nodes/model_loader/zmongo_model_introspector_node.py create mode 100644 nodes/model_loader/zmongo_universal_adapter_node.py create mode 100644 nodes/model_loader/zmongo_universal_model_loader_node.py create mode 100644 nodes/model_loader/zmongo_universal_model_selector_node.py create mode 100644 nodes/zmongo_field_selector_node.py create mode 100644 nodes/zmongo_flattened_field_selector_node.py create mode 100644 nodes/zmongo_nodes.py create mode 100644 nodes/zmongo_record_editor_node.py create mode 100644 nodes/zmongo_record_splitter.py create mode 100644 nodes/zmongo_tabular_record_view_node.py create mode 100644 web/js/app.js create mode 100644 web/js/preset_panel.js create mode 100644 web/js/zmongo_field_selector.js create mode 100644 web/js/zmongo_field_selector_dropdown.js create mode 100644 web/js/zmongo_record_editor.js create mode 100644 web/js/zmongo_tabular_record_view.js create mode 100644 web/js/zmongo_universal_adapter.js diff --git a/api/zmongo_record_editor_api.py b/api/zmongo_record_editor_api.py new file mode 100644 index 0000000..6fc7ee6 --- /dev/null +++ b/api/zmongo_record_editor_api.py @@ -0,0 +1,224 @@ +import json +import logging +from typing import Any, Dict, List + +from aiohttp import web +from bson import ObjectId + +from .zmongo_record_editor_node import ZMongoRecordEditorNode +from .zmongo_toolbag.zmongo import ZMongo +from .zmongo_toolbag.data_processor import DataProcessor + +logger = logging.getLogger(__name__) + + +def register_zmongo_record_editor_routes(prompt_server_instance): + if not prompt_server_instance: + logger.warning("PromptServer instance not available; record editor routes not registered.") + return + + routes = prompt_server_instance.routes + + def _clean_record_id(value: Any) -> str: + raw = "" if value is None else str(value) + cleaned = raw.strip() + + if len(cleaned) >= 2 and cleaned[0] == cleaned[-1] and cleaned[0] in {"'", '"'}: + cleaned = cleaned[1:-1].strip() + + return cleaned + + def _build_id_candidates(value: Any) -> List[Any]: + cleaned = _clean_record_id(value) + if not cleaned: + return [] + + candidates: List[Any] = [cleaned] + if ObjectId.is_valid(cleaned): + oid = ObjectId(cleaned) + candidates.insert(0, oid) + + return candidates + + def _parse_json_or_string(value: Any) -> Any: + raw = "" if value is None else str(value) + stripped = raw.strip() + if not stripped: + return "" + try: + return json.loads(stripped) + except Exception: + return raw + + def _result_succeeded(result: Any) -> bool: + return bool(result) and bool(getattr(result, "success", False)) + + def _result_data(result: Any) -> Dict[str, Any]: + data = getattr(result, "data", {}) or {} + return data if isinstance(data, dict) else {} + + @routes.get("/zmongo/record_editor/load") + async def zmongo_record_editor_load(request): + try: + collection_name = request.rel_url.query.get("collection_name", "").strip() + record_id = _clean_record_id(request.rel_url.query.get("record_id", "")) + selected_record_json = request.rel_url.query.get("selected_record_json", "") + + record = ZMongoRecordEditorNode.fetch_record( + collection_name=collection_name, + record_id=record_id, + fallback_record_json=selected_record_json, + ) + + if not record: + return web.json_response( + { + "success": True, + "collection_name": collection_name, + "record_id": record_id, + "record": {}, + "fields": [], + } + ) + + resolved_record_id = str(record.get("_id", record_id)) + + pairs = ZMongoRecordEditorNode.get_flattened_record_pairs( + collection_name=collection_name, + record_id=resolved_record_id, + fallback_record_json=selected_record_json, + ) + + fields = [{"path": path, "value": value} for path, value in pairs] + + return web.json_response( + { + "success": True, + "collection_name": collection_name, + "record_id": resolved_record_id, + "record": DataProcessor.to_json_compatible(record), + "fields": fields, + } + ) + except Exception as exc: + logger.exception("Failed to load record editor payload: %s", exc) + return web.json_response( + { + "success": False, + "error": str(exc), + "record": {}, + "fields": [], + }, + status=500, + ) + + @routes.post("/zmongo/record_editor/save") + async def zmongo_record_editor_save(request): + try: + payload: Dict[str, Any] = await request.json() + + collection_name = str(payload.get("collection_name", "")).strip() + record_id = _clean_record_id(payload.get("record_id", "")) + changes = payload.get("changes", {}) + + if not collection_name: + return web.json_response( + {"success": False, "error": "collection_name is required"}, + status=400, + ) + + if not record_id: + return web.json_response( + {"success": False, "error": "record_id is required"}, + status=400, + ) + + if not isinstance(changes, dict) or not changes: + return web.json_response( + {"success": False, "error": "changes must be a non-empty object"}, + status=400, + ) + + normalized_changes = { + str(path).strip(): _parse_json_or_string(value) + for path, value in changes.items() + if str(path).strip() + } + + if not normalized_changes: + return web.json_response( + {"success": False, "error": "No valid change paths were provided"}, + status=400, + ) + + zmongo = ZMongo() + update_doc = {"$set": normalized_changes} + + result = None + matched_any_candidate = False + + for candidate_id in _build_id_candidates(record_id): + result = zmongo.update_one( + collection_name, + {"_id": candidate_id}, + update_doc, + upsert=False, + ) + + if _result_succeeded(result): + data = _result_data(result) + matched_count = int(data.get("matched_count", 0) or 0) + modified_count = int(data.get("modified_count", 0) or 0) + upserted_id = data.get("upserted_id") + + if matched_count > 0: + matched_any_candidate = True + + if modified_count > 0 or upserted_id is not None or matched_count > 0: + break + + if result is None: + return web.json_response( + {"success": False, "error": "record_id is required"}, + status=400, + ) + + refreshed = ZMongoRecordEditorNode.fetch_record( + collection_name=collection_name, + record_id=record_id, + fallback_record_json="", + ) + + if not matched_any_candidate and not refreshed: + return web.json_response( + { + "success": False, + "error": f"No record found for _id '{record_id}' in collection '{collection_name}'", + "result": getattr(result, "to_dict", lambda: {})(), + "collection_name": collection_name, + "record_id": record_id, + "record": {}, + }, + status=404, + ) + + return web.json_response( + { + "success": bool(getattr(result, "success", False)), + "message": getattr(result, "message", "Saved"), + "error": getattr(result, "error", None), + "result": getattr(result, "to_dict", lambda: {})(), + "collection_name": collection_name, + "record_id": record_id, + "record": DataProcessor.to_json_compatible(refreshed), + } + ) + except Exception as exc: + logger.exception("Failed to save record editor changes: %s", exc) + return web.json_response( + { + "success": False, + "error": str(exc), + }, + status=500, + ) \ No newline at end of file diff --git a/api/zmongo_tabular_record_view_api.py b/api/zmongo_tabular_record_view_api.py new file mode 100644 index 0000000..4ad81ab --- /dev/null +++ b/api/zmongo_tabular_record_view_api.py @@ -0,0 +1,93 @@ +import logging +from aiohttp import web + +from .zmongo_tabular_record_view_node import ZMongoTabularRecordViewNode + +logger = logging.getLogger(__name__) + + +def register_zmongo_tabular_record_view_routes(prompt_server_instance): + if not prompt_server_instance: + logger.warning("PromptServer instance not available; tabular record routes not registered.") + return + + routes = prompt_server_instance.routes + + @routes.get("/zmongo/tabular_records") + async def zmongo_get_tabular_records(request): + try: + collection_name = request.rel_url.query.get("collection_name", "").strip() + + headings, flat_records, record_ids = ZMongoTabularRecordViewNode.get_table_payload(collection_name) + default_record_id = record_ids[0] if record_ids else "" + + return web.json_response( + { + "success": True, + "collection_name": collection_name, + "headings": headings, + "flat_records": flat_records, + "record_ids": record_ids, + "default_record_id": default_record_id, + "record_count": len(flat_records), + } + ) + except Exception as exc: + logger.exception("Failed to get tabular records: %s", exc) + return web.json_response( + { + "success": False, + "error": str(exc), + "headings": [], + "flat_records": [], + "record_ids": [], + "default_record_id": "", + "record_count": 0, + }, + status=500, + ) + + @routes.get("/zmongo/tabular_records_search") + async def zmongo_get_tabular_records_search(request): + try: + collection_name = request.rel_url.query.get("collection_name", "").strip() + search_text = request.rel_url.query.get("search_text", "") + flattened_field_name = request.rel_url.query.get("flattened_field_name", "").strip() + + headings, flat_records, record_ids = ( + ZMongoTabularRecordViewNode.get_filtered_table_payload( + collection_name=collection_name, + search_text=search_text, + flattened_field_name=flattened_field_name, + ) + ) + + default_record_id = record_ids[0] if record_ids else "" + + return web.json_response( + { + "success": True, + "collection_name": collection_name, + "search_text": search_text, + "flattened_field_name": flattened_field_name, + "headings": headings, + "flat_records": flat_records, + "record_ids": record_ids, + "default_record_id": default_record_id, + "record_count": len(flat_records), + } + ) + except Exception as exc: + logger.exception("Failed to search tabular records: %s", exc) + return web.json_response( + { + "success": False, + "error": str(exc), + "headings": [], + "flat_records": [], + "record_ids": [], + "default_record_id": "", + "record_count": 0, + }, + status=500, + ) \ No newline at end of file diff --git a/example_workflows/zmongo_db_browser.png b/examples/zmongo_chat_agent_starter.jpg similarity index 100% rename from example_workflows/zmongo_db_browser.png rename to examples/zmongo_chat_agent_starter.jpg diff --git a/example_workflows/zmongo_loop_sd15_demo.png b/examples/zmongo_db_browser.jpg similarity index 100% rename from example_workflows/zmongo_loop_sd15_demo.png rename to examples/zmongo_db_browser.jpg diff --git a/example_workflows/zmongo_record_loop_demoa - Copy.jpg b/examples/zmongo_loop_sd15_demo.jpg similarity index 100% rename from example_workflows/zmongo_record_loop_demoa - Copy.jpg rename to examples/zmongo_loop_sd15_demo.jpg diff --git a/example_workflows/zmongo_record_loop_demoa.png b/examples/zmongo_record_loop_demoa - Copy.jpg similarity index 100% rename from example_workflows/zmongo_record_loop_demoa.png rename to examples/zmongo_record_loop_demoa - Copy.jpg diff --git a/example_workflows/zmongo_record_loop_demob.jpg b/examples/zmongo_record_loop_demoa.jpg similarity index 100% rename from example_workflows/zmongo_record_loop_demob.jpg rename to examples/zmongo_record_loop_demoa.jpg diff --git a/example_workflows/zmongo_system_demoa.png b/examples/zmongo_system_demoa.jpg similarity index 100% rename from example_workflows/zmongo_system_demoa.png rename to examples/zmongo_system_demoa.jpg diff --git a/example_workflows/zmongo_system_demoa.json b/examples/zmongo_system_demoa.json similarity index 100% rename from example_workflows/zmongo_system_demoa.json rename to examples/zmongo_system_demoa.json diff --git a/nodes/model_loader/__init__.py b/nodes/model_loader/__init__.py new file mode 100644 index 0000000..9cb1109 --- /dev/null +++ b/nodes/model_loader/__init__.py @@ -0,0 +1,17 @@ +# SPECIALIZED MODEL NODE IMPORTS 🤖 +from .zmongo_model_compatibility_display_node import ZMongoModelCompatibilityDisplayNode +from .zmongo_model_introspector_node import ZMongoModelIntrospectorNode +from .zmongo_universal_model_selector_node import ZMongoUniversalModelSelectorNode +from .zmongo_universal_model_loader_node import ZMongoUniversalModelLoaderNode +from .zmongo_universal_adapter_node import ZMongoUniversalModelAdapterNode +from .zmongo_builtin_loader_adapter_node import ZMongoBuiltInLoaderAdapterNode + +# EXPORTING THE SUITE 📦 +__all__ = [ + "ZMongoModelCompatibilityDisplayNode", + "ZMongoModelIntrospectorNode", + "ZMongoUniversalModelSelectorNode", + "ZMongoUniversalModelLoaderNode", + "ZMongoUniversalModelAdapterNode", + "ZMongoBuiltInLoaderAdapterNode", +] \ No newline at end of file diff --git a/nodes/model_loader/zmongo_builtin_loader_adapter_node.py b/nodes/model_loader/zmongo_builtin_loader_adapter_node.py new file mode 100644 index 0000000..425bcc3 --- /dev/null +++ b/nodes/model_loader/zmongo_builtin_loader_adapter_node.py @@ -0,0 +1,437 @@ +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any, Dict, Optional, Tuple + +import folder_paths +from nodes import ( + CLIPLoader, + CLIPVisionLoader, + CheckpointLoaderSimple, + ControlNetLoader, + StyleModelLoader, + UNETLoader, + VAELoader, +) +from comfy_extras.nodes_upscale_model import UpscaleModelLoader + +class ZMongoBuiltInLoaderAdapterNode: + """ + Universal adapter that reuses ComfyUI built-in loaders wherever possible. + """ + + CATEGORY = "ZMongo/Models" + FUNCTION = "adapt" + + RETURN_TYPES = ( + "MODEL", + "CLIP", + "VAE", + "MODEL", + "CONTROL_NET", + "CLIP_VISION", + "UPSCALE_MODEL", + "STYLE_MODEL", + "STRING", + "BOOLEAN", + "STRING", + "STRING", + ) + RETURN_NAMES = ( + "model", + "clip", + "vae", + "unet_model", + "control_net", + "clip_vision", + "upscale_model", + "style_model", + "status_text", + "is_compatible", + "model_type", + "built_in_loader_name", + ) + + MODEL_TYPE_TO_FOLDER = { + "checkpoint": "checkpoints", + "vae": "vae", + "clip": "clip", + "text_encoder": "text_encoders", + "unet": "diffusion_models", + "diffusion": "diffusion_models", + "controlnet": "controlnet", + "clip_vision": "clip_vision", + "upscaler": "upscale_models", + "style_model": "style_models", + } + + BUILTIN_BY_MODEL_TYPE = { + "checkpoint": "CheckpointLoaderSimple", + "vae": "VAELoader", + "clip": "CLIPLoader", + "text_encoder": "CLIPLoader", + "unet": "UNETLoader", + "diffusion": "UNETLoader", + "controlnet": "ControlNetLoader", + "clip_vision": "CLIPVisionLoader", + "upscaler": "UpscaleModelLoader", + "style_model": "StyleModelLoader", + } + + SUPPORTED_MODEL_TYPES = set(BUILTIN_BY_MODEL_TYPE.keys()) + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "full_path": ("STRING", {"forceInput": True}), + }, + "optional": { + "model_info_json": ("STRING", {"forceInput": True}), + "clip_type": ( + [ + "stable_diffusion", + "stable_cascade", + "sd3", + "stable_audio", + "mochi", + "ltxv", + "pixart", + "cosmos", + "lumina2", + "wan", + "hidream", + ], + {"default": "stable_diffusion"}, + ), + }, + } + + @classmethod + def IS_CHANGED( + cls, + full_path: str, + model_info_json: str = "", + clip_type: str = "stable_diffusion", + ): + return full_path, model_info_json, clip_type + + @staticmethod + def _none_outputs() -> Tuple[Any, Any, Any, Any, Any, Any, Any]: + return None, None, None, None, None, None, None + + @staticmethod + def _safe_parse_model_info(model_info_json: str) -> Dict[str, Any]: + if not model_info_json: + return {} + try: + parsed = json.loads(model_info_json) + return parsed if isinstance(parsed, dict) else {} + except Exception: + return {} + + @staticmethod + def _normalize_path(full_path: str) -> Optional[Path]: + if not full_path: + return None + try: + return Path(full_path).expanduser().resolve() + except Exception: + return None + + @staticmethod + def _infer_model_type_from_path(path: Path) -> str: + parts = [part.lower() for part in path.parts] + + if "checkpoints" in parts: + return "checkpoint" + if "vae" in parts: + return "vae" + if "clip_vision" in parts: + return "clip_vision" + if "clip" in parts: + return "clip" + if "text_encoders" in parts: + return "text_encoder" + if "controlnet" in parts: + return "controlnet" + if "upscale_models" in parts: + return "upscaler" + if "style_models" in parts: + return "style_model" + if "diffusion_models" in parts: + return "diffusion" + if "unet" in parts: + return "unet" + + ext = path.suffix.lower() + if ext == ".ckpt": + return "checkpoint" + if ext in {".pt", ".pth", ".bin", ".safetensors"}: + return "unknown" + + return "unknown" + + @classmethod + def _resolve_model_type(cls, full_path: Path, model_info: Dict[str, Any]) -> str: + model_type = str(model_info.get("model_type", "")).strip() + if model_type: + return model_type + return cls._infer_model_type_from_path(full_path) + + @classmethod + def _relative_name_for_folder_type(cls, full_path: Path, folder_type: str) -> str: + normalized = full_path.resolve() + available = folder_paths.get_filename_list(folder_type) + + for entry in available: + candidate = Path(folder_paths.get_full_path(folder_type, entry)).resolve() + if candidate == normalized: + return entry.replace("\\", "/") + + raise ValueError( + f"Selected file is not registered under Comfy models folder type " + f"'{folder_type}': {full_path}" + ) + + @staticmethod + def _call_first_available(instance: Any, methods_with_kwargs): + last_error = None + + for method_name, kwargs in methods_with_kwargs: + method = getattr(instance, method_name, None) + if method is None: + continue + try: + return method(**kwargs) + except TypeError as exc: + last_error = exc + continue + + if last_error is not None: + raise last_error + + raise AttributeError( + f"No compatible loader method found on {instance.__class__.__name__}" + ) + + def _load_checkpoint(self, relative_name: str): + loader = CheckpointLoaderSimple() + result = self._call_first_available( + loader, + [ + ("load_checkpoint", {"ckpt_name": relative_name}), + ], + ) + if not isinstance(result, tuple) or len(result) < 3: + raise ValueError("CheckpointLoaderSimple returned an unexpected result.") + return result[0], result[1], result[2] + + def _load_vae(self, relative_name: str): + loader = VAELoader() + result = self._call_first_available( + loader, + [ + ("load_vae", {"vae_name": relative_name}), + ], + ) + if not isinstance(result, tuple) or len(result) < 1: + raise ValueError("VAELoader returned an unexpected result.") + return result[0] + + def _load_clip(self, relative_name: str, clip_type: str): + loader = CLIPLoader() + result = self._call_first_available( + loader, + [ + ("load_clip", {"clip_name": relative_name, "type": clip_type}), + ("load_clip", {"clip_name": relative_name}), + ], + ) + if not isinstance(result, tuple) or len(result) < 1: + raise ValueError("CLIPLoader returned an unexpected result.") + return result[0] + + def _load_unet(self, relative_name: str): + loader = UNETLoader() + result = self._call_first_available( + loader, + [ + ("load_unet", {"unet_name": relative_name}), + ("load_model", {"unet_name": relative_name}), + ], + ) + if not isinstance(result, tuple) or len(result) < 1: + raise ValueError("UNETLoader returned an unexpected result.") + return result[0] + + def _load_controlnet(self, relative_name: str): + loader = ControlNetLoader() + result = self._call_first_available( + loader, + [ + ("load_controlnet", {"control_net_name": relative_name}), + ("load_controlnet", {"controlnet_name": relative_name}), + ], + ) + if not isinstance(result, tuple) or len(result) < 1: + raise ValueError("ControlNetLoader returned an unexpected result.") + return result[0] + + def _load_clip_vision(self, relative_name: str): + loader = CLIPVisionLoader() + result = self._call_first_available( + loader, + [ + ("load_clip", {"clip_name": relative_name}), + ("load_clip_vision", {"clip_name": relative_name}), + ("load_clip_vision", {"clip_vision_name": relative_name}), + ], + ) + if not isinstance(result, tuple) or len(result) < 1: + raise ValueError("CLIPVisionLoader returned an unexpected result.") + return result[0] + + def _load_upscale_model(self, relative_name: str): + loader = UpscaleModelLoader() + result = self._call_first_available( + loader, + [ + ("load_model", {"model_name": relative_name}), + ("load_model", {"upscale_model": relative_name}), + ], + ) + if not isinstance(result, tuple) or len(result) < 1: + raise ValueError("UpscaleModelLoader returned an unexpected result.") + return result[0] + + def _load_style_model(self, relative_name: str): + loader = StyleModelLoader() + result = self._call_first_available( + loader, + [ + ("load_style_model", {"style_model_name": relative_name}), + ("load_model", {"style_model_name": relative_name}), + ("load_model", {"model_name": relative_name}), + ], + ) + if not isinstance(result, tuple) or len(result) < 1: + raise ValueError("StyleModelLoader returned an unexpected result.") + return result[0] + + def adapt( + self, + full_path: str, + model_info_json: str = "", + clip_type: str = "stable_diffusion", + ): + model_obj, clip_obj, vae_obj, unet_obj, controlnet_obj, clipvision_obj, upscale_obj = self._none_outputs() + style_obj = None + + path = self._normalize_path(full_path) + if path is None or not path.exists() or not path.is_file(): + return ( + model_obj, + clip_obj, + vae_obj, + unet_obj, + controlnet_obj, + clipvision_obj, + upscale_obj, + style_obj, + "Invalid or missing model path.", + False, + "unknown", + "", + ) + + model_info = self._safe_parse_model_info(model_info_json) + model_type = self._resolve_model_type(path, model_info) + built_in_loader_name = self.BUILTIN_BY_MODEL_TYPE.get(model_type, "") + + if model_type not in self.SUPPORTED_MODEL_TYPES: + return ( + model_obj, + clip_obj, + vae_obj, + unet_obj, + controlnet_obj, + clipvision_obj, + upscale_obj, + style_obj, + ( + f"No compatible built-in Comfy loader for model type " + f"'{model_type}'. Path: {path}" + ), + False, + model_type, + built_in_loader_name, + ) + + folder_type = self.MODEL_TYPE_TO_FOLDER[model_type] + + try: + relative_name = self._relative_name_for_folder_type(path, folder_type) + + if model_type == "checkpoint": + model_obj, clip_obj, vae_obj = self._load_checkpoint(relative_name) + elif model_type == "vae": + vae_obj = self._load_vae(relative_name) + elif model_type in {"clip", "text_encoder"}: + clip_obj = self._load_clip(relative_name, clip_type) + elif model_type in {"unet", "diffusion"}: + unet_obj = self._load_unet(relative_name) + elif model_type == "controlnet": + controlnet_obj = self._load_controlnet(relative_name) + elif model_type == "clip_vision": + clipvision_obj = self._load_clip_vision(relative_name) + elif model_type == "upscaler": + upscale_obj = self._load_upscale_model(relative_name) + elif model_type == "style_model": + style_obj = self._load_style_model(relative_name) + + status_text = ( + f"Loaded with built-in loader '{built_in_loader_name}'.\n" + f"Model type: {model_type}\n" + f"Folder type: {folder_type}\n" + f"Relative name: {relative_name}\n" + f"Path: {path}" + ) + + return ( + model_obj, + clip_obj, + vae_obj, + unet_obj, + controlnet_obj, + clipvision_obj, + upscale_obj, + style_obj, + status_text, + True, + model_type, + built_in_loader_name, + ) + + except Exception as exc: + return ( + model_obj, + clip_obj, + vae_obj, + unet_obj, + controlnet_obj, + clipvision_obj, + upscale_obj, + style_obj, + ( + f"Built-in loader adapter failed.\n" + f"Model type: {model_type}\n" + f"Built-in loader: {built_in_loader_name}\n" + f"Path: {path}\n" + f"Error: {exc}" + ), + False, + model_type, + built_in_loader_name, + ) \ No newline at end of file diff --git a/nodes/model_loader/zmongo_model_compatibility_display_node.py b/nodes/model_loader/zmongo_model_compatibility_display_node.py new file mode 100644 index 0000000..49d44e5 --- /dev/null +++ b/nodes/model_loader/zmongo_model_compatibility_display_node.py @@ -0,0 +1,116 @@ +from __future__ import annotations + +import json +from typing import Dict, List, Tuple + + +class ZMongoModelCompatibilityDisplayNode: + CATEGORY = "ZMongo/Models" + FUNCTION = "describe_compatibility" + RETURN_TYPES = ( + "BOOLEAN", # is_compatible + "STRING", # status_text + "STRING", # compatible_targets + "STRING", # compatible_sources + "STRING", # selected_model_type + "STRING", # selected_architecture + ) + RETURN_NAMES = ( + "is_compatible", + "status_text", + "compatible_targets", + "compatible_sources", + "selected_model_type", + "selected_architecture", + ) + + EXPECTED_NODE_TO_MODEL_TYPES: Dict[str, List[str]] = { + "ZMongoLLMLoaderNode": ["llm"], + "ZMongoPromptRunnerNode": ["llm"], + "ZMongoChatModelNode": ["llm"], + + "ZMongoEmbeddingLoaderNode": ["embedding", "text_encoder", "onnx_model"], + "ZMongoVectorIndexNode": ["embedding"], + "ZMongoEmbeddingInferenceNode": ["embedding", "text_encoder", "onnx_model"], + + "ZMongoCLIPLoaderNode": ["clip", "onnx_model"], + "ZMongoCLIPVisionLoaderNode": ["clip_vision", "onnx_model"], + "ZMongoVisionInferenceNode": ["clip", "clip_vision", "onnx_model", "vision"], + + "ZMongoCheckpointLoaderNode": ["checkpoint", "torch_model"], + "ZMongoTorchLoaderNode": ["torch_model", "checkpoint"], + "ZMongoTorchInferenceNode": ["torch_model", "checkpoint", "diffusion"], + + "ZMongoVAELoaderNode": ["vae"], + "ZMongoDiffusionLoaderNode": ["diffusion", "checkpoint"], + "ZMongoControlNetLoaderNode": ["controlnet"], + "ZMongoUpscaleLoaderNode": ["upscaler"], + "ZMongoLatentUpscaleLoaderNode": ["latent_upscaler"], + + "ZMongoAudioLoaderNode": ["audio_model"], + "ZMongoAudioInferenceNode": ["audio_model"], + "ZMongoSpeechLoaderNode": ["speech_model"], + "ZMongoSpeechInferenceNode": ["speech_model"], + } + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "model_info_json": ("STRING", {"forceInput": True}), + "expected_target_node": (sorted(list(cls.EXPECTED_NODE_TO_MODEL_TYPES.keys())),), + } + } + + @classmethod + def IS_CHANGED(cls, model_info_json: str, expected_target_node: str): + return model_info_json, expected_target_node + + @staticmethod + def _safe_load_model_info(model_info_json: str) -> dict: + try: + value = json.loads(model_info_json) + return value if isinstance(value, dict) else {} + except Exception: + return {} + + @staticmethod + def _csv(items: List[str]) -> str: + return ", ".join(items) if items else "None" + + def describe_compatibility(self, model_info_json: str, expected_target_node: str): + model_info = self._safe_load_model_info(model_info_json) + + selected_model_type = str(model_info.get("model_type", "unknown")) + selected_architecture = str(model_info.get("architecture_family", "unknown")) + compatible_targets = list(model_info.get("can_connect_to", [])) + compatible_sources = list(model_info.get("can_receive_from", [])) + + expected_types = self.EXPECTED_NODE_TO_MODEL_TYPES.get(expected_target_node, []) + is_compatible = selected_model_type in expected_types + + if is_compatible: + status_text = ( + f"Compatible: model_type '{selected_model_type}' can connect to " + f"'{expected_target_node}'." + ) + else: + status_text = ( + f"Incompatible: model_type '{selected_model_type}' does not match " + f"'{expected_target_node}'. Expected one of: {self._csv(expected_types)}." + ) + + status_text += ( + f"\nSelected architecture: {selected_architecture}" + f"\nCan connect to: {self._csv(compatible_targets)}" + f"\nCan receive from: {self._csv(compatible_sources)}" + ) + + return ( + is_compatible, + status_text, + self._csv(compatible_targets), + self._csv(compatible_sources), + selected_model_type, + selected_architecture, + ) \ No newline at end of file diff --git a/nodes/model_loader/zmongo_model_introspector_node.py b/nodes/model_loader/zmongo_model_introspector_node.py new file mode 100644 index 0000000..4364be5 --- /dev/null +++ b/nodes/model_loader/zmongo_model_introspector_node.py @@ -0,0 +1,395 @@ +from __future__ import annotations + +import json +from pathlib import Path +from typing import Dict, List, Tuple + + +class ZMongoModelIntrospectorNode: + CATEGORY = "ZMongo/Models" + FUNCTION = "inspect_model" + RETURN_TYPES = ( + "STRING", # model_info_json + "STRING", # summary_text + "STRING", # model_type + "STRING", # architecture_family + "STRING", # can_connect_to + "STRING", # can_receive_from + "STRING", # full_path + ) + RETURN_NAMES = ( + "model_info_json", + "summary_text", + "model_type", + "architecture_family", + "can_connect_to", + "can_receive_from", + "full_path", + ) + + FOLDER_TYPE_HINTS: Dict[str, Tuple[str, str]] = { + "llm": ("llm", "llama_cpp"), + "embeddings": ("embedding", "embedding_loader"), + "encoders": ("embedding", "embedding_loader"), + "clip": ("clip", "clip_loader"), + "clip_vision": ("clip_vision", "clip_vision_loader"), + "vae": ("vae", "vae_loader"), + "unet": ("unet", "torch_loader"), + "checkpoints": ("checkpoint", "checkpoint_loader"), + "diffusion_models": ("diffusion", "diffusion_loader"), + "loras": ("lora", "lora_loader"), + "controlnet": ("controlnet", "controlnet_loader"), + "upscale_models": ("upscaler", "upscale_loader"), + "latent_upscale_models": ("latent_upscaler", "latent_upscale_loader"), + "onnx": ("onnx_model", "onnxruntime"), + "vision": ("vision", "vision_loader"), + "audio": ("audio_model", "audio_loader"), + "speech_models": ("speech_model", "audio_loader"), + "text_encoders": ("text_encoder", "text_encoder_loader"), + } + + EXTENSION_HINTS: Dict[str, Tuple[str, str]] = { + ".gguf": ("llm", "llama_cpp"), + ".onnx": ("onnx_model", "onnxruntime"), + ".safetensors": ("torch_model", "torch_loader"), + ".ckpt": ("checkpoint", "checkpoint_loader"), + ".pt": ("torch_model", "torch_loader"), + ".pth": ("torch_model", "torch_loader"), + ".bin": ("binary_model", "custom_loader"), + ".json": ("config", "config_parser"), + ".yaml": ("config", "config_parser"), + ".yml": ("config", "config_parser"), + ".model": ("generic_model", "custom_loader"), + ".index": ("index", "index_loader"), + } + + ARCHITECTURE_HINTS = { + "bge": "bge", + "e5": "e5", + "gte": "gte", + "clip": "clip", + "siglip": "siglip", + "llama": "llama", + "mistral": "mistral", + "qwen": "qwen", + "gemma": "gemma", + "bert": "bert", + "t5": "t5", + "whisper": "whisper", + } + + COMPATIBILITY_MAP: Dict[str, Dict[str, List[str]]] = { + "llm": { + "can_connect_to": [ + "ZMongoLLMLoaderNode", + "ZMongoPromptRunnerNode", + "ZMongoChatModelNode", + ], + "can_receive_from": [ + "ZMongoUniversalModelSelectorNode", + "ZMongoModelTypeRouterNode", + ], + }, + "embedding": { + "can_connect_to": [ + "ZMongoEmbeddingLoaderNode", + "ZMongoVectorIndexNode", + "ZMongoEmbeddingInferenceNode", + ], + "can_receive_from": [ + "ZMongoUniversalModelSelectorNode", + "ZMongoModelTypeRouterNode", + ], + }, + "clip": { + "can_connect_to": [ + "ZMongoCLIPLoaderNode", + "ZMongoVisionInferenceNode", + ], + "can_receive_from": [ + "ZMongoUniversalModelSelectorNode", + "ZMongoModelTypeRouterNode", + ], + }, + "clip_vision": { + "can_connect_to": [ + "ZMongoCLIPVisionLoaderNode", + "ZMongoVisionInferenceNode", + ], + "can_receive_from": [ + "ZMongoUniversalModelSelectorNode", + "ZMongoModelTypeRouterNode", + ], + }, + "onnx_model": { + "can_connect_to": [ + "ZMongoONNXLoaderNode", + "ZMongoVisionInferenceNode", + "ZMongoEmbeddingInferenceNode", + ], + "can_receive_from": [ + "ZMongoUniversalModelSelectorNode", + "ZMongoModelTypeRouterNode", + ], + }, + "checkpoint": { + "can_connect_to": [ + "ZMongoCheckpointLoaderNode", + "ZMongoTorchInferenceNode", + ], + "can_receive_from": [ + "ZMongoUniversalModelSelectorNode", + "ZMongoModelTypeRouterNode", + ], + }, + "torch_model": { + "can_connect_to": [ + "ZMongoTorchLoaderNode", + "ZMongoTorchInferenceNode", + ], + "can_receive_from": [ + "ZMongoUniversalModelSelectorNode", + "ZMongoModelTypeRouterNode", + ], + }, + "text_encoder": { + "can_connect_to": [ + "ZMongoTextEncoderLoaderNode", + "ZMongoEmbeddingInferenceNode", + ], + "can_receive_from": [ + "ZMongoUniversalModelSelectorNode", + "ZMongoModelTypeRouterNode", + ], + }, + "vae": { + "can_connect_to": [ + "ZMongoVAELoaderNode", + ], + "can_receive_from": [ + "ZMongoUniversalModelSelectorNode", + "ZMongoModelTypeRouterNode", + ], + }, + "diffusion": { + "can_connect_to": [ + "ZMongoDiffusionLoaderNode", + ], + "can_receive_from": [ + "ZMongoUniversalModelSelectorNode", + "ZMongoModelTypeRouterNode", + ], + }, + "controlnet": { + "can_connect_to": [ + "ZMongoControlNetLoaderNode", + ], + "can_receive_from": [ + "ZMongoUniversalModelSelectorNode", + "ZMongoModelTypeRouterNode", + ], + }, + "upscaler": { + "can_connect_to": [ + "ZMongoUpscaleLoaderNode", + ], + "can_receive_from": [ + "ZMongoUniversalModelSelectorNode", + "ZMongoModelTypeRouterNode", + ], + }, + "latent_upscaler": { + "can_connect_to": [ + "ZMongoLatentUpscaleLoaderNode", + ], + "can_receive_from": [ + "ZMongoUniversalModelSelectorNode", + "ZMongoModelTypeRouterNode", + ], + }, + "audio_model": { + "can_connect_to": [ + "ZMongoAudioLoaderNode", + "ZMongoAudioInferenceNode", + ], + "can_receive_from": [ + "ZMongoUniversalModelSelectorNode", + "ZMongoModelTypeRouterNode", + ], + }, + "speech_model": { + "can_connect_to": [ + "ZMongoSpeechLoaderNode", + "ZMongoSpeechInferenceNode", + ], + "can_receive_from": [ + "ZMongoUniversalModelSelectorNode", + "ZMongoModelTypeRouterNode", + ], + }, + "unknown": { + "can_connect_to": [], + "can_receive_from": [ + "ZMongoUniversalModelSelectorNode", + "ZMongoModelTypeRouterNode", + ], + }, + } + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "full_path": ("STRING", {"forceInput": True}), + } + } + + @classmethod + def IS_CHANGED(cls, full_path: str): + return full_path + + @staticmethod + def _normalize_path(full_path: str) -> Path | None: + if not full_path: + return None + try: + return Path(full_path).expanduser().resolve() + except Exception: + return None + + @classmethod + def _folder_hint(cls, path: Path) -> Tuple[str, str] | None: + for part in reversed([p.lower() for p in path.parts]): + if part in cls.FOLDER_TYPE_HINTS: + return cls.FOLDER_TYPE_HINTS[part] + return None + + @classmethod + def _extension_hint(cls, path: Path) -> Tuple[str, str] | None: + return cls.EXTENSION_HINTS.get(path.suffix.lower()) + + @classmethod + def _architecture_family(cls, path: Path) -> str: + haystack = " ".join([part.lower() for part in path.parts]) + for token, family in cls.ARCHITECTURE_HINTS.items(): + if token in haystack: + return family + return "unknown" + + @classmethod + def _classify(cls, path: Path) -> Tuple[str, str]: + ext = path.suffix.lower() + folder_hint = cls._folder_hint(path) + ext_hint = cls._extension_hint(path) + + if ext == ".gguf": + return "llm", "llama_cpp" + if ext == ".onnx": + return "onnx_model", "onnxruntime" + if ext in {".json", ".yaml", ".yml"}: + if folder_hint: + return f"config_for_{folder_hint[0]}", folder_hint[1] + return "config", "config_parser" + if ext in {".safetensors", ".pt", ".pth", ".ckpt", ".bin", ".model"}: + if folder_hint: + return folder_hint + if ext_hint: + return ext_hint + if folder_hint: + return folder_hint + if ext_hint: + return ext_hint + return "unknown", "manual_loader_required" + + @staticmethod + def _find_sidecar_files(path: Path) -> List[str]: + candidates = [ + "config.json", + "tokenizer.json", + "tokenizer_config.json", + "model_index.json", + "params.json", + "metadata.json", + "config.yaml", + "config.yml", + ] + found: List[str] = [] + for candidate in candidates: + candidate_path = path.parent / candidate + if candidate_path.exists() and candidate_path.is_file(): + found.append(str(candidate_path)) + return found + + @staticmethod + def _format_list(items: List[str]) -> str: + return ", ".join(items) if items else "None" + + def inspect_model(self, full_path: str): + path = self._normalize_path(full_path) + if path is None or not path.exists() or not path.is_file(): + model_info = { + "full_path": "", + "model_type": "unknown", + "loader_hint": "invalid_path", + "architecture_family": "unknown", + "filename": "", + "extension": "", + "parent_folder": "", + "sidecar_files": [], + "can_connect_to": [], + "can_receive_from": [], + } + summary = "Invalid or missing model path." + return ( + json.dumps(model_info, indent=2), + summary, + "unknown", + "unknown", + "", + "", + "", + ) + + model_type, loader_hint = self._classify(path) + architecture_family = self._architecture_family(path) + compatibility = self.COMPATIBILITY_MAP.get(model_type, self.COMPATIBILITY_MAP["unknown"]) + can_connect_to = compatibility["can_connect_to"] + can_receive_from = compatibility["can_receive_from"] + sidecar_files = self._find_sidecar_files(path) + + model_info = { + "full_path": str(path), + "model_type": model_type, + "loader_hint": loader_hint, + "architecture_family": architecture_family, + "filename": path.name, + "stem": path.stem, + "extension": path.suffix.lower(), + "parent_folder": path.parent.name, + "sidecar_files": sidecar_files, + "can_connect_to": can_connect_to, + "can_receive_from": can_receive_from, + } + + summary_lines = [ + f"Model: {path.name}", + f"Type: {model_type}", + f"Architecture: {architecture_family}", + f"Loader hint: {loader_hint}", + f"Extension: {path.suffix.lower()}", + f"Folder: {path.parent.name}", + f"Can connect to: {self._format_list(can_connect_to)}", + f"Can receive from: {self._format_list(can_receive_from)}", + f"Sidecar files: {self._format_list(sidecar_files)}", + ] + summary_text = "\n".join(summary_lines) + + return ( + json.dumps(model_info, indent=2), + summary_text, + model_type, + architecture_family, + self._format_list(can_connect_to), + self._format_list(can_receive_from), + str(path), + ) \ No newline at end of file diff --git a/nodes/model_loader/zmongo_universal_adapter_node.py b/nodes/model_loader/zmongo_universal_adapter_node.py new file mode 100644 index 0000000..307b5be --- /dev/null +++ b/nodes/model_loader/zmongo_universal_adapter_node.py @@ -0,0 +1,222 @@ +from __future__ import annotations + +import json +from typing import Any, Dict, List, Tuple + + +class ZMongoUniversalModelAdapterNode: + """ + Universal adapter that consumes ZMongo model metadata and exposes a fixed + set of outputs. It uses smart-check compatibility metadata to decide which + outputs are active. + + Important: + - Output socket types remain statically declared for ComfyUI. + - Inactive outputs are returned as None and should be treated as disabled. + - A companion JS extension can visually dim inactive sockets/widgets. + """ + + CATEGORY = "ZMongo/Models" + FUNCTION = "adapt" + + RETURN_TYPES = ( + "MODEL", + "CLIP", + "VAE", + "STRING", + "STRING", + "BOOLEAN", + "STRING", + ) + RETURN_NAMES = ( + "model", + "clip", + "vae", + "status_text", + "active_outputs_json", + "is_compatible", + "model_type", + ) + + # Which outputs are allowed for each detected model type. + OUTPUTS_BY_MODEL_TYPE: Dict[str, List[str]] = { + "checkpoint": ["model", "clip", "vae"], + "torch_model": ["model"], + "diffusion": ["model"], + "controlnet": ["model"], + "upscaler": ["model"], + "latent_upscaler": ["model"], + "unet": ["model"], + "clip": ["clip"], + "clip_vision": ["clip"], + "vae": ["vae"], + "llm": [], + "embedding": [], + "text_encoder": [], + "onnx_model": [], + "audio_model": [], + "speech_model": [], + "unknown": [], + } + + # Which target families the selected model is allowed to drive. + TARGETS_BY_MODEL_TYPE: Dict[str, List[str]] = { + "checkpoint": ["MODEL", "CLIP", "VAE"], + "torch_model": ["MODEL"], + "diffusion": ["MODEL"], + "controlnet": ["MODEL"], + "upscaler": ["MODEL"], + "latent_upscaler": ["MODEL"], + "unet": ["MODEL"], + "clip": ["CLIP"], + "clip_vision": ["CLIP"], + "vae": ["VAE"], + "llm": [], + "embedding": [], + "text_encoder": [], + "onnx_model": [], + "audio_model": [], + "speech_model": [], + "unknown": [], + } + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "zmongo_model_json": ("STRING", {"forceInput": True}), + }, + "optional": { + "expected_output_family": ( + ["AUTO", "MODEL", "CLIP", "VAE"], + {"default": "AUTO"}, + ), + }, + "hidden": { + "unique_id": "UNIQUE_ID", + }, + } + + @classmethod + def IS_CHANGED(cls, zmongo_model_json: str, expected_output_family: str = "AUTO", unique_id=None): + return (zmongo_model_json, expected_output_family, unique_id) + + @classmethod + def VALIDATE_INPUTS(cls, input_types): + # Keep backend validation permissive here; real validation happens after + # parsing the metadata payload because the smart-check logic depends on + # model metadata, not just the declared socket type. + return True + + @staticmethod + def _safe_parse(payload: str) -> Dict[str, Any]: + try: + data = json.loads(payload) + return data if isinstance(data, dict) else {} + except Exception: + return {} + + @classmethod + def _compute_active_outputs( + cls, + model_type: str, + expected_output_family: str, + ) -> Tuple[List[str], bool, str]: + allowed = cls.OUTPUTS_BY_MODEL_TYPE.get(model_type, []) + allowed_families = cls.TARGETS_BY_MODEL_TYPE.get(model_type, []) + + if expected_output_family == "AUTO": + return allowed, bool(allowed), "AUTO mode" + + is_compatible = expected_output_family in allowed_families + if not is_compatible: + return [], False, ( + f"Model type '{model_type}' is not compatible with requested " + f"output family '{expected_output_family}'. Allowed: " + f"{', '.join(allowed_families) if allowed_families else 'None'}" + ) + + # Keep only the outputs that map to the requested family. + family_to_outputs = { + "MODEL": ["model"], + "CLIP": ["clip"], + "VAE": ["vae"], + } + requested_outputs = family_to_outputs.get(expected_output_family, []) + active = [name for name in allowed if name in requested_outputs] + return active, True, f"Compatible with requested family '{expected_output_family}'" + + @staticmethod + def _status_text( + model_type: str, + architecture: str, + full_path: str, + active_outputs: List[str], + compatible: bool, + reason: str, + ) -> str: + return "\n".join( + [ + f"Model type: {model_type}", + f"Architecture: {architecture}", + f"Compatible: {compatible}", + f"Active outputs: {', '.join(active_outputs) if active_outputs else 'None'}", + f"Reason: {reason}", + f"Path: {full_path}", + ] + ) + + def adapt( + self, + zmongo_model_json: str, + expected_output_family: str = "AUTO", + unique_id: str | None = None, + ): + del unique_id + + payload = self._safe_parse(zmongo_model_json) + model_type = str(payload.get("model_type", "unknown")) + architecture = str(payload.get("architecture_family", "unknown")) + full_path = str(payload.get("full_path", "")) + + active_outputs, is_compatible, reason = self._compute_active_outputs( + model_type=model_type, + expected_output_family=expected_output_family, + ) + + # Deliberately conservative: + # We do not fabricate real Comfy MODEL/CLIP/VAE objects here. + # Those should be created by specialized loader/adapter nodes. + model_obj = None + clip_obj = None + vae_obj = None + + status_text = self._status_text( + model_type=model_type, + architecture=architecture, + full_path=full_path, + active_outputs=active_outputs, + compatible=is_compatible, + reason=reason, + ) + + active_outputs_json = json.dumps( + { + "active_outputs": active_outputs, + "expected_output_family": expected_output_family, + "model_type": model_type, + "architecture_family": architecture, + "compatible": is_compatible, + }, + indent=2, + ) + + return ( + model_obj if "model" in active_outputs else None, + clip_obj if "clip" in active_outputs else None, + vae_obj if "vae" in active_outputs else None, + status_text, + active_outputs_json, + is_compatible, + model_type, + ) \ No newline at end of file diff --git a/nodes/model_loader/zmongo_universal_model_loader_node.py b/nodes/model_loader/zmongo_universal_model_loader_node.py new file mode 100644 index 0000000..c28a223 --- /dev/null +++ b/nodes/model_loader/zmongo_universal_model_loader_node.py @@ -0,0 +1,378 @@ +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any, Dict, List, Optional, Tuple + + +class ZMongoUniversalModelLoaderNode: + """ + Universal model loader facade for ZMongo. + + This node does not try to fully instantiate every possible model backend. + Instead, it: + + 1. Accepts a selected model path. + 2. Classifies the model. + 3. Returns a universal internal payload (JSON string). + 4. Activates only the outputs that make sense for the selected model. + 5. Exposes status / compatibility information for downstream nodes. + + Important: + - Comfy output socket types are static. + - So this node declares a fixed set of outputs and populates only the valid ones. + - Core Comfy outputs such as MODEL / CLIP / VAE are only populated when the + selected file appears compatible with those families. + """ + + CATEGORY = "ZMongo/Models" + FUNCTION = "load_model" + + RETURN_TYPES = ( + "STRING", # zmongo_model_json + "MODEL", # model + "CLIP", # clip + "VAE", # vae + "STRING", # model_info_json + "STRING", # status_text + "BOOLEAN", # is_valid + "STRING", # model_type + "STRING", # allowed_outputs + ) + RETURN_NAMES = ( + "zmongo_model_json", + "model", + "clip", + "vae", + "model_info_json", + "status_text", + "is_valid", + "model_type", + "allowed_outputs", + ) + + FOLDER_TYPE_HINTS: Dict[str, Tuple[str, str]] = { + "llm": ("llm", "llama_cpp"), + "embeddings": ("embedding", "embedding_loader"), + "encoders": ("embedding", "embedding_loader"), + "clip": ("clip", "clip_loader"), + "clip_vision": ("clip_vision", "clip_vision_loader"), + "vae": ("vae", "vae_loader"), + "unet": ("unet", "torch_loader"), + "checkpoints": ("checkpoint", "checkpoint_loader"), + "diffusion_models": ("diffusion", "diffusion_loader"), + "loras": ("lora", "lora_loader"), + "controlnet": ("controlnet", "controlnet_loader"), + "upscale_models": ("upscaler", "upscale_loader"), + "latent_upscale_models": ("latent_upscaler", "latent_upscale_loader"), + "onnx": ("onnx_model", "onnxruntime"), + "vision": ("vision", "vision_loader"), + "audio": ("audio_model", "audio_loader"), + "speech_models": ("speech_model", "audio_loader"), + "text_encoders": ("text_encoder", "text_encoder_loader"), + } + + EXTENSION_HINTS: Dict[str, Tuple[str, str]] = { + ".gguf": ("llm", "llama_cpp"), + ".onnx": ("onnx_model", "onnxruntime"), + ".safetensors": ("torch_model", "torch_loader"), + ".ckpt": ("checkpoint", "checkpoint_loader"), + ".pt": ("torch_model", "torch_loader"), + ".pth": ("torch_model", "torch_loader"), + ".bin": ("binary_model", "custom_loader"), + ".json": ("config", "config_parser"), + ".yaml": ("config", "config_parser"), + ".yml": ("config", "config_parser"), + ".model": ("generic_model", "custom_loader"), + ".index": ("index", "index_loader"), + } + + ARCHITECTURE_HINTS = { + "bge": "bge", + "e5": "e5", + "gte": "gte", + "clip": "clip", + "siglip": "siglip", + "llama": "llama", + "mistral": "mistral", + "qwen": "qwen", + "gemma": "gemma", + "bert": "bert", + "t5": "t5", + "whisper": "whisper", + "sd": "stable_diffusion", + "flux": "flux", + } + + MODEL_OUTPUT_COMPATIBILITY: Dict[str, List[str]] = { + "llm": ["zmongo_model_json"], + "embedding": ["zmongo_model_json"], + "text_encoder": ["zmongo_model_json"], + "onnx_model": ["zmongo_model_json"], + "audio_model": ["zmongo_model_json"], + "speech_model": ["zmongo_model_json"], + "clip": ["zmongo_model_json", "clip"], + "clip_vision": ["zmongo_model_json", "clip"], + "vae": ["zmongo_model_json", "vae"], + "checkpoint": ["zmongo_model_json", "model", "clip", "vae"], + "torch_model": ["zmongo_model_json", "model"], + "diffusion": ["zmongo_model_json", "model"], + "controlnet": ["zmongo_model_json", "model"], + "upscaler": ["zmongo_model_json", "model"], + "latent_upscaler": ["zmongo_model_json", "model"], + "unet": ["zmongo_model_json", "model"], + } + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "full_path": ("STRING", {"forceInput": True}), + }, + "optional": { + "model_info_json": ("STRING", {"forceInput": True}), + }, + } + + @classmethod + def IS_CHANGED(cls, full_path: str, model_info_json: str = ""): + return full_path, model_info_json + + @staticmethod + def _normalize_path(full_path: str) -> Optional[Path]: + if not full_path: + return None + try: + return Path(full_path).expanduser().resolve() + except Exception: + return None + + @classmethod + def _folder_hint(cls, path: Path) -> Optional[Tuple[str, str]]: + for part in reversed([p.lower() for p in path.parts]): + if part in cls.FOLDER_TYPE_HINTS: + return cls.FOLDER_TYPE_HINTS[part] + return None + + @classmethod + def _extension_hint(cls, path: Path) -> Optional[Tuple[str, str]]: + return cls.EXTENSION_HINTS.get(path.suffix.lower()) + + @classmethod + def _architecture_family(cls, path: Path) -> str: + haystack = " ".join([part.lower() for part in path.parts]) + for token, family in cls.ARCHITECTURE_HINTS.items(): + if token in haystack: + return family + return "unknown" + + @classmethod + def _classify(cls, path: Path) -> Tuple[str, str]: + ext = path.suffix.lower() + folder_hint = cls._folder_hint(path) + ext_hint = cls._extension_hint(path) + + if ext == ".gguf": + return "llm", "llama_cpp" + if ext == ".onnx": + return "onnx_model", "onnxruntime" + if ext in {".json", ".yaml", ".yml"}: + if folder_hint: + return f"config_for_{folder_hint[0]}", folder_hint[1] + return "config", "config_parser" + if ext in {".safetensors", ".pt", ".pth", ".ckpt", ".bin", ".model"}: + if folder_hint: + return folder_hint + if ext_hint: + return ext_hint + if folder_hint: + return folder_hint + if ext_hint: + return ext_hint + return "unknown", "manual_loader_required" + + @staticmethod + def _find_sidecar_files(path: Path) -> List[str]: + candidates = [ + "config.json", + "tokenizer.json", + "tokenizer_config.json", + "model_index.json", + "params.json", + "metadata.json", + "config.yaml", + "config.yml", + ] + found: List[str] = [] + for candidate in candidates: + candidate_path = path.parent / candidate + if candidate_path.exists() and candidate_path.is_file(): + found.append(str(candidate_path)) + return found + + @staticmethod + def _try_parse_model_info(model_info_json: str) -> Dict[str, Any]: + if not model_info_json: + return {} + try: + parsed = json.loads(model_info_json) + return parsed if isinstance(parsed, dict) else {} + except Exception: + return {} + + @classmethod + def _build_model_info(cls, path: Path) -> Dict[str, Any]: + model_type, loader_hint = cls._classify(path) + architecture_family = cls._architecture_family(path) + allowed_outputs = cls.MODEL_OUTPUT_COMPATIBILITY.get(model_type, ["zmongo_model_json"]) + + return { + "full_path": str(path), + "filename": path.name, + "stem": path.stem, + "extension": path.suffix.lower(), + "parent_folder": path.parent.name, + "model_type": model_type, + "loader_hint": loader_hint, + "architecture_family": architecture_family, + "sidecar_files": cls._find_sidecar_files(path), + "allowed_outputs": allowed_outputs, + "exists": path.exists(), + "is_file": path.is_file(), + } + + @staticmethod + def _make_status(model_info: Dict[str, Any]) -> str: + allowed_outputs = model_info.get("allowed_outputs", []) + return "\n".join( + [ + f"Loaded: {model_info.get('filename', '')}", + f"Type: {model_info.get('model_type', 'unknown')}", + f"Architecture: {model_info.get('architecture_family', 'unknown')}", + f"Loader hint: {model_info.get('loader_hint', 'unknown')}", + f"Allowed outputs: {', '.join(allowed_outputs) if allowed_outputs else 'None'}", + f"Path: {model_info.get('full_path', '')}", + ] + ) + + @staticmethod + def _empty_core_outputs() -> Tuple[Any, Any, Any]: + return None, None, None + + def _build_zmongo_model_payload(self, model_info: Dict[str, Any]) -> str: + payload = { + "kind": "zmongo_model", + "full_path": model_info.get("full_path", ""), + "model_type": model_info.get("model_type", "unknown"), + "architecture_family": model_info.get("architecture_family", "unknown"), + "loader_hint": model_info.get("loader_hint", "unknown"), + "allowed_outputs": model_info.get("allowed_outputs", []), + "sidecar_files": model_info.get("sidecar_files", []), + } + return json.dumps(payload, indent=2) + + def _activate_core_outputs( + self, + model_info: Dict[str, Any], + ) -> Tuple[Any, Any, Any]: + """ + Returns placeholders for MODEL / CLIP / VAE-compatible cases. + + Note: + This node intentionally does not instantiate arbitrary backend objects + for every model family. For now: + - checkpoint-like families can mark MODEL/CLIP/VAE as logically enabled + - actual object creation should be delegated to adapter/loader nodes later + """ + model_type = model_info.get("model_type", "unknown") + allowed_outputs = set(model_info.get("allowed_outputs", [])) + + model_obj = None + clip_obj = None + vae_obj = None + + # Deliberately conservative: + # We only advertise activation through metadata/status here. + # Returning real core objects should be done in a dedicated adapter node. + if model_type == "checkpoint": + model_obj = None + clip_obj = None + vae_obj = None + elif "model" in allowed_outputs: + model_obj = None + elif "clip" in allowed_outputs: + clip_obj = None + elif "vae" in allowed_outputs: + vae_obj = None + + return model_obj, clip_obj, vae_obj + + def load_model(self, full_path: str, model_info_json: str = ""): + parsed_path = self._normalize_path(full_path) + if parsed_path is None or not parsed_path.exists() or not parsed_path.is_file(): + empty_payload = json.dumps( + { + "kind": "zmongo_model", + "full_path": "", + "model_type": "unknown", + "architecture_family": "unknown", + "loader_hint": "invalid_path", + "allowed_outputs": [], + "sidecar_files": [], + }, + indent=2, + ) + empty_info = json.dumps( + { + "full_path": "", + "model_type": "unknown", + "loader_hint": "invalid_path", + "architecture_family": "unknown", + "allowed_outputs": [], + }, + indent=2, + ) + return ( + empty_payload, + None, + None, + None, + empty_info, + "Invalid or missing model path.", + False, + "unknown", + "", + ) + + provided_info = self._try_parse_model_info(model_info_json) + if provided_info and provided_info.get("full_path") == str(parsed_path): + model_info = provided_info + model_info.setdefault("allowed_outputs", self.MODEL_OUTPUT_COMPATIBILITY.get( + model_info.get("model_type", "unknown"), + ["zmongo_model_json"], + )) + model_info.setdefault("architecture_family", self._architecture_family(parsed_path)) + model_info.setdefault("loader_hint", self._classify(parsed_path)[1]) + model_info.setdefault("sidecar_files", self._find_sidecar_files(parsed_path)) + else: + model_info = self._build_model_info(parsed_path) + + zmongo_model_json = self._build_zmongo_model_payload(model_info) + model_obj, clip_obj, vae_obj = self._activate_core_outputs(model_info) + + status_text = self._make_status(model_info) + is_valid = model_info.get("model_type", "unknown") != "unknown" + model_type = str(model_info.get("model_type", "unknown")) + allowed_outputs = ", ".join(model_info.get("allowed_outputs", [])) + + return ( + zmongo_model_json, + model_obj, + clip_obj, + vae_obj, + json.dumps(model_info, indent=2), + status_text, + is_valid, + model_type, + allowed_outputs, + ) \ No newline at end of file diff --git a/nodes/model_loader/zmongo_universal_model_selector_node.py b/nodes/model_loader/zmongo_universal_model_selector_node.py new file mode 100644 index 0000000..879c34a --- /dev/null +++ b/nodes/model_loader/zmongo_universal_model_selector_node.py @@ -0,0 +1,133 @@ +from __future__ import annotations + +from pathlib import Path +from typing import Iterable, List, Tuple + +import folder_paths + + +class ZMongoUniversalModelSelectorNode: + """ + Universal recursive model selector for files stored under: + + ComfyUI/models/zmongo/ + + It does not load the model itself. It only returns normalized path data so + downstream ZMongo nodes can decide how to use the selected file. + """ + + CATEGORY = "ZMongo/Models" + FUNCTION = "select_model" + RETURN_TYPES = ("STRING", "STRING", "STRING", "STRING", "STRING") + RETURN_NAMES = ("full_path", "relative_path", "filename", "stem", "extension") + + # File extensions commonly used by Comfy / ML / embeddings / LLM ecosystems. + ALLOWED_EXTENSIONS = { + ".bin", + ".ckpt", + ".gguf", + ".index", + ".json", + ".mar", + ".model", + ".onnx", + ".pt", + ".pth", + ".py", + ".safetensors", + ".t7", + ".tar", + ".uf2", + ".yaml", + ".yml", + } + + @classmethod + def _zmongo_root(cls) -> Path: + return Path(folder_paths.models_dir) / "checkpoints" + + @classmethod + def _iter_model_files(cls, root: Path) -> Iterable[Path]: + if not root.exists() or not root.is_dir(): + return [] + + files: List[Path] = [] + for path in root.rglob("*"): + if not path.is_file(): + continue + if path.suffix.lower() not in cls.ALLOWED_EXTENSIONS: + continue + files.append(path) + + files.sort(key=lambda p: str(p.relative_to(root)).lower()) + return files + + @classmethod + def _relative_choices(cls) -> List[str]: + root = cls._zmongo_root() + files = cls._iter_model_files(root) + + choices = [str(path.relative_to(root)).replace("\\", "/") for path in files] + + if not choices: + choices = [""] + + return choices + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "model_relative_path": (cls._relative_choices(),), + }, + "optional": { + "refresh_token": ( + "INT", + { + "default": 0, + "min": 0, + "max": 999999, + "step": 1, + "display": "number", + }, + ), + }, + } + + @classmethod + def IS_CHANGED(cls, model_relative_path: str, refresh_token: int = 0): + """ + Re-executes when selection changes or when refresh_token changes. + This does not rebuild the combo live in the UI, but it helps rerun nodes + after you change the filesystem and bump the token. + """ + return (model_relative_path, refresh_token) + + @staticmethod + def _empty_result() -> Tuple[str, str, str, str, str]: + return ("", "", "", "", "") + + def select_model( + self, + model_relative_path: str, + refresh_token: int = 0, + ) -> Tuple[str, str, str, str, str]: + del refresh_token # execution-only knob for Comfy cache / rerun behavior + + root = self._zmongo_root() + + if model_relative_path == "": + return self._empty_result() + + full_path = (root / model_relative_path).resolve() + + if not full_path.exists() or not full_path.is_file(): + return self._empty_result() + + return ( + str(full_path), + model_relative_path.replace("\\", "/"), + full_path.name, + full_path.stem, + full_path.suffix.lower(), + ) \ No newline at end of file diff --git a/nodes/zmongo_field_selector_node.py b/nodes/zmongo_field_selector_node.py new file mode 100644 index 0000000..72f1a64 --- /dev/null +++ b/nodes/zmongo_field_selector_node.py @@ -0,0 +1,213 @@ +import json +import logging +from typing import Any, Dict, List + +logger = logging.getLogger(__name__) +from ..zmongo_toolbag.zmongo import ZMongo +from ..zmongo_toolbag.data_processor import DataProcessor + +class ZMongoFieldSelectorNode: + CATEGORY = "ZMongo" + FUNCTION = "select_field" + RETURN_TYPES = ("STRING", "STRING", "INT", "STRING", "STRING") + RETURN_NAMES = ( + "collection_name", + "selected_field_name", + "selected_field_index", + "selected_field_value", + "flattened_field_names_json", + ) + + @classmethod + def _get_zmongo(cls) -> ZMongo: + return ZMongo() + + @staticmethod + def _clamp_index(value: int, max_len: int) -> int: + if max_len <= 0: + return 0 + return max(0, min(int(value), max_len - 1)) + + @staticmethod + def _stringify_value(value: Any) -> str: + if value is None: + return "" + if isinstance(value, str): + return value + if isinstance(value, (dict, list, tuple, bool, int, float)): + try: + return json.dumps( + DataProcessor.to_json_compatible(value), + ensure_ascii=False, + indent=2, + default=str, + ) + except Exception: + return str(value) + return str(value) + + @classmethod + def get_collection_names(cls) -> List[str]: + zmongo = None + try: + zmongo = cls._get_zmongo() + result = zmongo.run_sync(zmongo.list_collections_async) + + if not result or not getattr(result, "success", False): + logger.warning( + "Could not load collection names: %s", + getattr(result, "error", None), + ) + return [""] + + data = result.data or {} + if not isinstance(data, dict): + return [""] + + names = data.get("collections", []) + if not isinstance(names, list): + return [""] + + names = [str(name) for name in names if str(name).strip()] + return names or [""] + + except Exception as exc: + logger.exception("Error loading collection names: %s", exc) + return [""] + finally: + if zmongo: + try: + zmongo.close() + except Exception: + pass + + @classmethod + def get_sample_document(cls, collection_name: str) -> Dict[str, Any]: + zmongo = None + try: + if not collection_name or collection_name.startswith("<"): + return {} + + zmongo = cls._get_zmongo() + result = zmongo.find_one(collection_name, {}) + + if not result or not getattr(result, "success", False): + logger.warning( + "Could not load sample document for '%s': %s", + collection_name, + getattr(result, "error", None), + ) + return {} + + doc = result.original() if hasattr(result, "original") else result.data + return doc if isinstance(doc, dict) else {} + + except Exception as exc: + logger.exception( + "Error loading sample document for collection '%s': %s", + collection_name, + exc, + ) + return {} + finally: + if zmongo: + try: + + + zmongo.close() + except Exception: + pass + + @classmethod + def get_flattened_field_names(cls, collection_name: str) -> List[str]: + try: + doc = cls.get_sample_document(collection_name) + if not doc: + return [] + + flattened = DataProcessor.flatten_json(doc) + if not isinstance(flattened, dict): + return [] + + return sorted(str(k) for k in flattened.keys()) + except Exception as exc: + logger.exception( + "Error flattening fields for collection '%s': %s", + collection_name, + exc, + ) + return [] + + @classmethod + def INPUT_TYPES(cls) -> Dict[str, Any]: + collections = cls.get_collection_names() + default_collection = collections[0] if collections else "" + + return { + "required": { + "collection_name": (collections, {"default": default_collection}), + "field_name": ([""], {"default": ""}), + "field_index": ( + "INT", + { + "default": 0, + "min": 0, + "max": 100000, + "step": 1, + "display": "number", + }, + ), + } + } + + @classmethod + def VALIDATE_INPUTS(cls, collection_name, field_name, field_index): + return True + + @classmethod + def IS_CHANGED(cls, collection_name: str, field_name: str, field_index: int): + return f"{collection_name}|{field_name}|{field_index}" + + def select_field(self, collection_name: str, field_name: str, field_index: int): + doc = self.get_sample_document(collection_name) + field_names = self.get_flattened_field_names(collection_name) + + if not doc or not field_names: + return ( + str(collection_name), + "", + 0, + "", + json.dumps([], indent=2), + ) + + normalized_index = self._clamp_index(field_index, len(field_names)) + + selected_field_name = field_name + if ( + not selected_field_name + or selected_field_name.startswith("<") + or selected_field_name not in field_names + ): + selected_field_name = field_names[normalized_index] + + selected_index = field_names.index(selected_field_name) + selected_value = DataProcessor.get_value(doc, selected_field_name) + selected_value_str = self._stringify_value(selected_value) + + return ( + str(collection_name), + str(selected_field_name), + int(selected_index), + selected_value_str, + json.dumps(field_names, indent=2), + ) + + +NODE_CLASS_MAPPINGS = { + "ZMongoFlattenedFieldDropdownNode": ZMongoFieldSelectorNode, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "ZMongoFlattenedFieldDropdownNode": "🧩 ZMongo Flattened Field Dropdown Node", +} \ No newline at end of file diff --git a/nodes/zmongo_flattened_field_selector_node.py b/nodes/zmongo_flattened_field_selector_node.py new file mode 100644 index 0000000..88965ca --- /dev/null +++ b/nodes/zmongo_flattened_field_selector_node.py @@ -0,0 +1,213 @@ +import json +import logging +from typing import Any, Dict, List + +from .zmongo_toolbag.zmongo import ZMongo +from .zmongo_toolbag.data_processor import DataProcessor + +logger = logging.getLogger(__name__) + + +class ZMongoFlattenedFieldSelectorNode: + CATEGORY = "ZMongo" + FUNCTION = "select_field" + RETURN_TYPES = ("STRING", "STRING", "INT", "STRING", "STRING") + RETURN_NAMES = ( + "collection_name", + "selected_field_name", + "selected_field_index", + "selected_field_value", + "flattened_field_names_json", + ) + + @classmethod + def _get_zmongo(cls) -> ZMongo: + return ZMongo() + + @staticmethod + def _clamp_index(value: int, max_len: int) -> int: + if max_len <= 0: + return 0 + return max(0, min(int(value), max_len - 1)) + + @staticmethod + def _stringify_value(value: Any) -> str: + if value is None: + return "" + if isinstance(value, str): + return value + if isinstance(value, (dict, list, tuple, bool, int, float)): + try: + return json.dumps( + DataProcessor.to_json_compatible(value), + ensure_ascii=False, + indent=2, + default=str, + ) + except Exception: + return str(value) + return str(value) + + @classmethod + def get_collection_names(cls) -> List[str]: + zmongo = None + try: + zmongo = cls._get_zmongo() + result = zmongo.run_sync(zmongo.list_collections_async) + + if not result or not getattr(result, "success", False): + logger.warning( + "Could not load collection names: %s", + getattr(result, "error", None), + ) + return [""] + + data = result.data or {} + if not isinstance(data, dict): + return [""] + + names = data.get("collections", []) + if not isinstance(names, list): + return [""] + + names = [str(name) for name in names if str(name).strip()] + return names or [""] + + except Exception as exc: + logger.exception("Error loading collection names: %s", exc) + return [""] + finally: + if zmongo: + try: + zmongo.close() + except Exception: + pass + + @classmethod + def get_sample_document(cls, collection_name: str) -> Dict[str, Any]: + zmongo = None + try: + if not collection_name or collection_name.startswith("<"): + return {} + + zmongo = cls._get_zmongo() + result = zmongo.find_one(collection_name, {}) + + if not result or not getattr(result, "success", False): + logger.warning( + "Could not load sample document for '%s': %s", + collection_name, + getattr(result, "error", None), + ) + return {} + + doc = result.original() if hasattr(result, "original") else result.data + return doc if isinstance(doc, dict) else {} + + except Exception as exc: + logger.exception( + "Error loading sample document for collection '%s': %s", + collection_name, + exc, + ) + return {} + finally: + if zmongo: + try: + zmongo.close() + except Exception: + pass + + @classmethod + def get_flattened_field_names(cls, collection_name: str) -> List[str]: + try: + doc = cls.get_sample_document(collection_name) + if not doc: + return [] + + flattened = DataProcessor.flatten_json(doc) + if not isinstance(flattened, dict): + return [] + + return sorted(str(k) for k in flattened.keys()) + except Exception as exc: + logger.exception( + "Error flattening fields for collection '%s': %s", + collection_name, + exc, + ) + return [] + + @classmethod + def INPUT_TYPES(cls) -> Dict[str, Any]: + collections = cls.get_collection_names() + default_collection = collections[0] if collections else "" + + return { + "required": { + "collection_name": (collections, {"default": default_collection}), + "field_name": ([""], {"default": ""}), + "field_index": ( + "INT", + { + "default": 0, + "min": 0, + "max": 100000, + "step": 1, + "display": "number", + }, + ), + } + } + + @classmethod + def VALIDATE_INPUTS(cls, collection_name, field_name, field_index): + return True + + @classmethod + def IS_CHANGED(cls, collection_name: str, field_name: str, field_index: int): + return f"{collection_name}|{field_name}|{field_index}" + + def select_field(self, collection_name: str, field_name: str, field_index: int): + doc = self.get_sample_document(collection_name) + field_names = self.get_flattened_field_names(collection_name) + + if not doc or not field_names: + return ( + str(collection_name), + "", + 0, + "", + json.dumps([], indent=2), + ) + + normalized_index = self._clamp_index(field_index, len(field_names)) + + selected_field_name = field_name + if ( + not selected_field_name + or selected_field_name.startswith("<") + or selected_field_name not in field_names + ): + selected_field_name = field_names[normalized_index] + + selected_index = field_names.index(selected_field_name) + selected_value = DataProcessor.get_value(doc, selected_field_name) + selected_value_str = self._stringify_value(selected_value) + + return ( + str(collection_name), + str(selected_field_name), + int(selected_index), + selected_value_str, + json.dumps(field_names, indent=2), + ) + + +NODE_CLASS_MAPPINGS = { + "ZMongoFlattenedFieldSelectorNode": ZMongoFlattenedFieldSelectorNode, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "ZMongoFlattenedFieldSelectorNode": "🧩 ZMongo Flattened Field Selector Node", +} \ No newline at end of file diff --git a/nodes/zmongo_nodes.py b/nodes/zmongo_nodes.py new file mode 100644 index 0000000..edf3a94 --- /dev/null +++ b/nodes/zmongo_nodes.py @@ -0,0 +1,2788 @@ +import json +import logging +import random +import threading +from typing import Any, Dict, List, Optional + +import requests +from bson import json_util +from bson.objectid import ObjectId +from pymongo import MongoClient +from pymongo.errors import DuplicateKeyError +from ..zmongo_toolbag.safe_result import SafeResult +from ..zmongo_toolbag.zmongo import ZMongo +from ..zmongo_toolbag.zembedder import ZEmbedder +from ..zmongo_toolbag.data_processor import DataProcessor + +logger = logging.getLogger(__name__) + + +_CONFIG_LOCK = threading.Lock() +_ZMONGO_SINGLETON: Optional[ZMongo] = None +_CURRENT_URI: str = "mongodb://127.0.0.1:27017" +_CURRENT_DB: str = "test" + +from ..zmongo_toolbag.data_processor import DataProcessor + +def _safe_json(obj: Any) -> str: + try: + if isinstance(obj, SafeResult): + return DataProcessor.to_json(obj.to_dict(), indent=2) + return DataProcessor.to_json(obj, indent=2) + except Exception: + try: + return json.dumps( + json.loads(json_util.dumps(obj)), + indent=2, + ensure_ascii=False, + default=str, + ) + except Exception as exc: + return json.dumps({"error": f"Serialization failed: {exc}"}, indent=2) + + +def _coerce_records(data_json: str) -> SafeResult: + if not data_json or not data_json.strip(): + return SafeResult.fail("No data provided") + + try: + data = json.loads(data_json) + except Exception as exc: + return SafeResult.fail(f"Invalid JSON: {exc}") + + if isinstance(data, list): + return SafeResult.ok(data) + if isinstance(data, dict): + return SafeResult.ok([data]) + return SafeResult.fail("JSON payload must be an object or an array of objects") + + +def _build_summary_from_hits(hits: List[Dict[str, Any]]) -> str: + if not hits: + return "No matching documents found." + + lines: List[str] = [] + for index, hit in enumerate(hits, start=1): + doc = hit.get("document") or {} + score = hit.get("retrieval_score", 0.0) + doc_id = doc.get("_id", "") + title = ( + doc.get("title") + or doc.get("name") + or doc.get("username") + or doc.get("email") + or doc.get("text") + or doc.get("content") + or "" + ) + if isinstance(title, str): + title = title.strip().replace("\n", " ") + if len(title) > 100: + title = title[:100] + "..." + else: + title = str(title) + lines.append(f"{index}. score={float(score):.4f} | _id={doc_id} | {title}") + return "\n".join(lines) + + +def _get_zmongo(uri: Optional[str] = None, db_name: Optional[str] = None) -> ZMongo: + global _ZMONGO_SINGLETON, _CURRENT_URI, _CURRENT_DB + + wanted_uri = uri or _CURRENT_URI + wanted_db = db_name or _CURRENT_DB + + with _CONFIG_LOCK: + needs_new = ( + _ZMONGO_SINGLETON is None + or _CURRENT_URI != wanted_uri + or _CURRENT_DB != wanted_db + ) + if needs_new: + old_instance = _ZMONGO_SINGLETON + _CURRENT_URI = wanted_uri + _CURRENT_DB = wanted_db + _ZMONGO_SINGLETON = ZMongo(uri=wanted_uri, db_name=wanted_db) + if old_instance is not None: + try: + old_instance.close() + except Exception: + logger.debug("Failed to close prior ZMongo instance", exc_info=True) + return _ZMONGO_SINGLETON + + +class ZMongoAPIMixin: + """Shared API/local access helpers for ZMongo nodes.""" + + API_COMMON_INPUTS = { + "zmongo_api_url": ("STRING", {"default": ""}), + "zmongo_api_key": ("STRING", {"default": "", "multiline": False}), + "use_api": ("BOOLEAN", {"default": False}), + "api_timeout_sec": ("INT", {"default": 10, "min": 1, "max": 120}), + } + + @staticmethod + def _parse_json_object(raw: str, field_name: str) -> Dict[str, Any]: + text = (raw or "").strip() + if not text: + return {} + try: + parsed = json.loads(text) + except Exception as exc: + raise ValueError(f"{field_name} is not valid JSON: {exc}") from exc + if not isinstance(parsed, dict): + raise ValueError(f"{field_name} must be a JSON object") + return parsed + + @staticmethod + def _normalize_sort(sort_field: str, sort_direction: str): + direction = 1 if sort_direction == "ascending" else -1 + field = (sort_field or "").strip() or "_id" + return [(field, direction)] + + @staticmethod + def _build_auth_headers(api_key: str) -> Dict[str, str]: + headers = {"Content-Type": "application/json"} + if api_key: + headers["Authorization"] = f"Bearer {api_key}" + return headers + + def _post_api(self, api_url: str, api_key: str, path: str, payload: Dict[str, Any], timeout_sec: int) -> Dict[str, Any]: + if not api_url or not api_url.strip(): + raise ValueError("zmongo_api_url is required when use_api=True") + + url = api_url.rstrip("/") + path + headers = self._build_auth_headers(api_key) + + response = requests.post(url, json=payload, headers=headers, timeout=int(timeout_sec)) + if response.status_code != 200: + raise RuntimeError(f"API request failed: {response.status_code} {response.text}") + + data = response.json() + if not data.get("success", False): + raise RuntimeError(f"API error: {data}") + return data + + def _fetch_record_by_index_api( + self, + api_url: str, + api_key: str, + database_name: str, + collection_name: str, + select_index: int, + sort_field: str, + sort_direction: str, + mongo_filter: Dict[str, Any], + projection: Dict[str, Any], + timeout_sec: int, + ) -> Dict[str, Any]: + payload = { + "database_name": database_name, + "collection_name": collection_name, + "index": int(select_index), + "sort": { + "field": sort_field or "_id", + "direction": sort_direction, + }, + "filter": mongo_filter or {}, + "projection": projection or {}, + } + return self._post_api( + api_url=api_url, + api_key=api_key, + path="/record/by_index", + payload=payload, + timeout_sec=timeout_sec, + ) + + def _fetch_records_range_api( + self, + api_url: str, + api_key: str, + database_name: str, + collection_name: str, + start_index: int, + end_index: int, + step: int, + sort_field: str, + sort_direction: str, + mongo_filter: Dict[str, Any], + projection: Dict[str, Any], + timeout_sec: int, + ) -> Dict[str, Any]: + payload = { + "database_name": database_name, + "collection_name": collection_name, + "start_index": int(start_index), + "end_index": int(end_index), + "step": int(step), + "sort": { + "field": sort_field or "_id", + "direction": sort_direction, + }, + "filter": mongo_filter or {}, + "projection": projection or {}, + } + return self._post_api( + api_url=api_url, + api_key=api_key, + path="/records/range", + payload=payload, + timeout_sec=timeout_sec, + ) + + def _fetch_collection_preview_api( + self, + api_url: str, + api_key: str, + database_name: str, + collection_name: str, + limit: int, + select_index: int, + timeout_sec: int, + ) -> Dict[str, Any]: + payload = { + "database_name": database_name, + "collection_name": collection_name, + "limit": int(limit), + "select_index": int(select_index), + } + return self._post_api( + api_url=api_url, + api_key=api_key, + path="/collection/preview", + payload=payload, + timeout_sec=timeout_sec, + ) + + +class ZMongoFlattenedFieldSelector: + """ + ComfyUI node that: + 1. Loads MongoDB collection names into a dropdown + 2. Loads one sample document from the selected collection + 3. Flattens that document with DataProcessor.flatten_json(...) + 4. Uses an INT input as the selected flattened field index + 5. Outputs the selected field name and related metadata + """ + + CATEGORY = "ZMongo" + RETURN_TYPES = ("STRING", "STRING", "INT", "STRING") + RETURN_NAMES = ( + "collection_name", + "selected_field_name", + "selected_field_index", + "flattened_field_names_json", + ) + FUNCTION = "select_field" + + @classmethod + def _get_zmongo(cls) -> ZMongo: + return ZMongo() + + @classmethod + def _safe_collection_names(cls) -> List[str]: + zmongo = None + try: + zmongo = cls._get_zmongo() + + # Prefer sync access through run_sync because current ZMongo.list_collections() + # is async and not a normal sync wrapper. + result = zmongo.run_sync(zmongo.list_collections_async) + + if not result or not getattr(result, "success", False): + logger.warning("Failed to load Mongo collections: %s", getattr(result, "error", None)) + return [""] + + data = result.data or {} + collections = data.get("collections", []) if isinstance(data, dict) else [] + + if not collections: + return [""] + + return [str(name) for name in collections] + + except Exception as exc: + logger.exception("Error retrieving collection names: %s", exc) + return [""] + finally: + if zmongo: + try: + zmongo.close() + except Exception: + pass + + @classmethod + def INPUT_TYPES(cls) -> Dict[str, Any]: + collections = cls._safe_collection_names() + + return { + "required": { + "collection_name": (collections,), + "field_index": ( + "INT", + { + "default": 0, + "min": 0, + "max": 100000, + "step": 1, + "display": "number", + }, + ), + } + } + + @classmethod + def IS_CHANGED(cls, collection_name: str, field_index: int): + # Causes reevaluation when either changes. + return f"{collection_name}:{field_index}" + + def _get_flattened_field_names(self, collection_name: str) -> List[str]: + zmongo = None + try: + if not collection_name or collection_name.startswith("<"): + return [] + + zmongo = self._get_zmongo() + + # Pull one sample document from the collection. + result = zmongo.find_one(collection_name, {}) + + if not result or not getattr(result, "success", False): + logger.warning( + "find_one failed for collection '%s': %s", + collection_name, + getattr(result, "error", None), + ) + return [] + + doc = result.original() if hasattr(result, "original") else result.data + if not isinstance(doc, dict) or not doc: + return [] + + flattened = DataProcessor.flatten_json(doc) + if not isinstance(flattened, dict): + return [] + + return list(flattened.keys()) + + except Exception as exc: + logger.exception("Error flattening fields for '%s': %s", collection_name, exc) + return [] + finally: + if zmongo: + try: + zmongo.close() + except Exception: + pass + + def select_field(self, collection_name: str, field_index: int): + field_names = self._get_flattened_field_names(collection_name) + + if not field_names: + return ( + str(collection_name), + "", + 0, + json.dumps([], indent=2), + ) + + normalized_index = max(0, min(int(field_index), len(field_names) - 1)) + selected_field_name = field_names[normalized_index] + + return ( + str(collection_name), + str(selected_field_name), + normalized_index, + json.dumps(field_names, indent=2), + ) + + +class ZMongoSaveValueNode: + """ + Saves a passed value from prior nodes into MongoDB using a flattened + dot-path field key. + + Supports: + - source_record_json as one object + - source_record_json as a list of objects + - choosing a record from the list with source_record_index + - explicit query override + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "mongo_uri": ("STRING", {"default": "mongodb://127.0.0.1:27017"}), + "database_name": ("STRING", {"default": "test"}), + + "source_record_json": ("STRING", {"forceInput": True}), + "source_record_index": ("INT", {"default": 1, "min": 1, "max": 1000000}), + "value_to_save": ("STRING", {"forceInput": True}), + + "target_collection": ("STRING", {"default": "comfy"}), + "target_field_path": ("STRING", {"default": "responses.output.text"}), + + "parse_value_as_json": ("BOOLEAN", {"default": False}), + "upsert_if_missing": ("BOOLEAN", {"default": False}), + + "explicit_query_json": ("STRING", {"default": "", "multiline": True}), + } + } + + RETURN_TYPES = ( + "STRING", # result_json + "STRING", # target_query_json + "STRING", # saved_value_json + "STRING", # target_collection_out + "STRING", # target_field_path_out + ) + RETURN_NAMES = ( + "result_json", + "target_query_json", + "saved_value_json", + "target_collection_out", + "target_field_path_out", + ) + FUNCTION = "save_value" + CATEGORY = "ZMongo/Operations" + + @staticmethod + def _parse_json_object(raw: str, field_name: str) -> Dict[str, Any]: + text = (raw or "").strip() + if not text: + return {} + try: + parsed = json.loads(text) + except Exception as exc: + raise ValueError(f"{field_name} is not valid JSON: {exc}") from exc + if not isinstance(parsed, dict): + raise ValueError(f"{field_name} must be a JSON object") + return parsed + + @staticmethod + def _parse_source_record(raw: str, source_record_index: int) -> Dict[str, Any]: + text = (raw or "").strip() + if not text: + raise ValueError("source_record_json is empty") + + try: + parsed = json.loads(text) + except Exception as exc: + raise ValueError(f"source_record_json is not valid JSON: {exc}") from exc + + if isinstance(parsed, dict): + return parsed + + if isinstance(parsed, list): + if not parsed: + raise ValueError("source_record_json list is empty") + + bounded_index = max(1, min(int(source_record_index), len(parsed))) + selected = parsed[bounded_index - 1] + + if not isinstance(selected, dict): + raise ValueError("Selected item from source_record_json list is not a JSON object") + + return selected + + raise ValueError("source_record_json must be a JSON object or a JSON list of objects") + + @staticmethod + def _parse_value(raw_value: Any, parse_value_as_json: bool) -> Any: + if raw_value is None: + return None + + if not parse_value_as_json: + return raw_value + + if isinstance(raw_value, str): + stripped = raw_value.strip() + if not stripped: + return "" + try: + return json.loads(stripped) + except Exception: + return raw_value + + return raw_value + + @staticmethod + def _extract_target_query( + source_record: Dict[str, Any], + explicit_query_json: str, + ) -> Dict[str, Any]: + explicit_query = ZMongoSaveValueNode._parse_json_object( + explicit_query_json, "explicit_query_json" + ) + if explicit_query: + return explicit_query + + if isinstance(source_record, dict) and source_record.get("_id") is not None: + return {"_id": source_record["_id"]} + + raise ValueError( + "No target query could be determined. Provide source_record_json with _id " + "or set explicit_query_json." + ) + + @staticmethod + def _extract_target_collection( + source_record: Dict[str, Any], + target_collection: str, + ) -> str: + if target_collection and target_collection.strip(): + return target_collection.strip() + + inferred = "" + if isinstance(source_record, dict): + inferred = str(source_record.get("_collection_name") or "").strip() + + if inferred: + return inferred + + raise ValueError( + "No target collection provided. Set target_collection or pass source_record_json " + "containing _collection_name." + ) + + def save_value( + self, + mongo_uri, + database_name, + source_record_json, + source_record_index, + value_to_save, + target_collection, + target_field_path, + parse_value_as_json, + upsert_if_missing, + explicit_query_json, + ): + try: + source_record = self._parse_source_record(source_record_json, source_record_index) + resolved_collection = self._extract_target_collection(source_record, target_collection) + target_query = self._extract_target_query(source_record, explicit_query_json) + + if not target_field_path or not str(target_field_path).strip(): + raise ValueError("target_field_path is required") + + resolved_field_path = str(target_field_path).strip() + parsed_value = self._parse_value(value_to_save, parse_value_as_json) + + update_doc = { + resolved_field_path: parsed_value + } + + zmongo = _get_zmongo(uri=mongo_uri, db_name=database_name) + result = zmongo.run_sync( + zmongo.update_one, + resolved_collection, + target_query, + update_doc, + upsert_if_missing, + ) + + if not isinstance(result, SafeResult): + result = SafeResult.fail( + f"Unexpected update result type: {type(result).__name__}" + ) + + result_payload = { + "success": result.success, + "message": result.message, + "error": result.error, + "data": result.data, + "target_collection": resolved_collection, + "target_field_path": resolved_field_path, + "target_query": target_query, + "source_record_index_used": int(source_record_index), + } + + return ( + _safe_json(result_payload), + _safe_json(target_query), + _safe_json(parsed_value), + str(resolved_collection), + str(resolved_field_path), + ) + + except Exception as exc: + error_payload = { + "success": False, + "error": str(exc), + "target_collection": target_collection, + "target_field_path": target_field_path, + } + return ( + json.dumps(error_payload, indent=2), + "{}", + _safe_json(value_to_save), + str(target_collection or ""), + str(target_field_path or ""), + ) + +class ZMongoSaveBatchTextNode: + """ + Saves a batch of text items to file(s) in the ComfyUI output directory. + + Intended use: + - true ComfyUI list input + - JSON list string input + - one file per item + - one combined file + - append or overwrite behavior + - stable indexed filenames for batch exports + """ + + INPUT_IS_LIST = True + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "batch_content": ("STRING", {"forceInput": True}), + "base_filename": ("STRING", {"default": "zmongo_batch"}), + "starting_index": ("INT", {"default": 1, "min": 0, "max": 1000000}), + "file_type": (["txt", "json", "md"],), + "save_behavior": ([ + "write new content", + "append to existing file", + ],), + "save_layout": ([ + "separate file per item", + "one combined file", + ],), + "reuse_exact_filenames": ("BOOLEAN", {"default": False}), + "skip_empty_items": ("BOOLEAN", {"default": True}), + "trim_whitespace": ("BOOLEAN", {"default": True}), + "join_separator": ("STRING", {"default": "\n"}), + } + } + + RETURN_TYPES = ("STRING", "INT", "STRING") + RETURN_NAMES = ( + "saved_file_path_or_paths", + "saved_file_count", + "summary", + ) + FUNCTION = "save_batch" + CATEGORY = "ZMongo/Output" + + @staticmethod + def _unwrap_scalar(value, default=None): + """ + ComfyUI may pass scalar inputs as one-item lists when INPUT_IS_LIST = True. + """ + if isinstance(value, list): + if not value: + return default + return value[0] + return value + + @classmethod + def _unwrap_and_cast_int(cls, value, default=0) -> int: + value = cls._unwrap_scalar(value, default) + if value is None: + return int(default) + return int(value) + + @classmethod + def _unwrap_and_cast_bool(cls, value, default=False) -> bool: + value = cls._unwrap_scalar(value, default) + if isinstance(value, bool): + return value + return bool(value) + + @classmethod + def _unwrap_and_cast_str(cls, value, default="") -> str: + value = cls._unwrap_scalar(value, default) + if value is None: + return str(default) + return str(value) + + @staticmethod + def _coerce_to_items(batch_content): + """ + Accepts: + - true ComfyUI list input + - JSON list string + - single plain string + """ + if isinstance(batch_content, list): + if len(batch_content) == 1 and isinstance(batch_content[0], str): + raw = batch_content[0].strip() + if raw: + try: + parsed = json.loads(raw) + if isinstance(parsed, list): + return parsed + except Exception: + pass + return batch_content + + if batch_content is None: + return [""] + + raw = str(batch_content) + stripped = raw.strip() + if not stripped: + return [""] + + try: + parsed = json.loads(stripped) + if isinstance(parsed, list): + return parsed + return [raw] + except Exception: + return [raw] + + @staticmethod + def _item_to_text(item, file_type: str, trim_whitespace: bool) -> str: + if item is None: + text = "" + elif isinstance(item, str): + text = item + elif file_type == "json": + text = json.dumps(item, indent=2, ensure_ascii=False) + elif isinstance(item, (dict, list)): + text = json.dumps(item, indent=2, ensure_ascii=False) + else: + text = str(item) + + if trim_whitespace and isinstance(text, str): + text = text.strip() + + return text + + @staticmethod + def _ensure_parent_dir(path: str): + import os + parent = os.path.dirname(path) + if parent: + os.makedirs(parent, exist_ok=True) + + @staticmethod + def _make_single_filename(base_filename: str, file_type: str) -> str: + return f"{base_filename}.{file_type}" + + @staticmethod + def _make_indexed_filename(base_filename: str, index: int, file_type: str) -> str: + return f"{base_filename}_{int(index):04d}.{file_type}" + + @staticmethod + def _make_unique_path(path: str, file_type: str) -> str: + import os + + if not os.path.exists(path): + return path + + base = path[: -(len(file_type) + 1)] + counter = 1 + while True: + candidate = f"{base}_{counter}.{file_type}" + if not os.path.exists(candidate): + return candidate + counter += 1 + + def save_batch( + self, + batch_content, + base_filename, + starting_index, + file_type, + save_behavior, + save_layout, + reuse_exact_filenames, + skip_empty_items, + trim_whitespace, + join_separator, + ): + try: + import os + import folder_paths + + # Normalize scalar inputs that may arrive as [value] + base_filename = self._unwrap_and_cast_str(base_filename, "zmongo_batch") + starting_index = self._unwrap_and_cast_int(starting_index, 1) + file_type = self._unwrap_and_cast_str(file_type, "txt") + save_behavior = self._unwrap_and_cast_str(save_behavior, "write new content") + save_layout = self._unwrap_and_cast_str(save_layout, "separate file per item") + reuse_exact_filenames = self._unwrap_and_cast_bool(reuse_exact_filenames, False) + skip_empty_items = self._unwrap_and_cast_bool(skip_empty_items, True) + trim_whitespace = self._unwrap_and_cast_bool(trim_whitespace, True) + join_separator = self._unwrap_and_cast_str(join_separator, "\n") + + output_dir = folder_paths.get_output_directory() + raw_items = self._coerce_to_items(batch_content) + + normalized_items = [] + for item in raw_items: + text_value = self._item_to_text(item, file_type, trim_whitespace) + if skip_empty_items and not text_value: + continue + normalized_items.append(text_value) + + if not normalized_items: + return ("[]", 0, "No non-empty items to save.") + + append_mode = (save_behavior == "append to existing file") + file_open_mode = "a" if append_mode else "w" + + saved_paths = [] + + if save_layout == "one combined file": + filename = self._make_single_filename(base_filename, file_type) + full_path = os.path.join(output_dir, filename) + + if not reuse_exact_filenames and not append_mode: + full_path = self._make_unique_path(full_path, file_type) + + self._ensure_parent_dir(full_path) + + with open(full_path, file_open_mode, encoding="utf-8") as f: + f.write(join_separator.join(normalized_items)) + + saved_paths.append(full_path) + + else: + for offset, item_text in enumerate(normalized_items): + item_index = int(starting_index) + offset + filename = self._make_indexed_filename(base_filename, item_index, file_type) + full_path = os.path.join(output_dir, filename) + + if not reuse_exact_filenames and not append_mode: + full_path = self._make_unique_path(full_path, file_type) + + self._ensure_parent_dir(full_path) + + with open(full_path, file_open_mode, encoding="utf-8") as f: + f.write(item_text) + + saved_paths.append(full_path) + + if len(saved_paths) == 1: + saved_output = saved_paths[0] + else: + saved_output = json.dumps(saved_paths, indent=2, ensure_ascii=False) + + summary = ( + f"Saved {len(saved_paths)} file(s)\n" + f"Base filename: {base_filename}\n" + f"File type: {file_type}\n" + f"Save behavior: {save_behavior}\n" + f"Save layout: {save_layout}" + ) + + return (saved_output, len(saved_paths), summary) + + except Exception as exc: + return (f"Error saving batch file(s): {exc}", 0, f"Error: {exc}") + + +class ZMongoDataPassThroughNode: + """ + Normalizes ZMongo output into formats that ComfyUI core nodes can consume. + + Main purposes: + - Pass through one text value + - Convert JSON list text into ComfyUI list output + - Convert JSON dict/list items into JSON strings + - Optionally split plain text by lines + - Provide both a single preview value and a list output for batch/list mapping + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "input_data": ("STRING", {"forceInput": True}), + "input_interpretation": ([ + "auto detect", + "treat as plain text", + "treat as json", + "treat as newline list", + ],), + "output_mode": ([ + "single item only", + "list output only", + "single item and list output", + ],), + "item_conversion": ([ + "keep strings as-is", + "convert everything to text", + "convert dicts and lists to json text", + ],), + "skip_empty_items": ("BOOLEAN", {"default": True}), + "trim_whitespace": ("BOOLEAN", {"default": True}), + } + } + + RETURN_TYPES = ( + "STRING", # single_item + "STRING", # list_items + "INT", # item_count + "STRING", # detected_input_type + "STRING", # normalized_json + ) + RETURN_NAMES = ( + "single_item", + "list_items", + "item_count", + "detected_input_type", + "normalized_json", + ) + OUTPUT_IS_LIST = (False, True, False, False, False) + FUNCTION = "pass_through" + CATEGORY = "ZMongo/Utility" + + @staticmethod + def _normalize_item(item: Any, item_conversion: str, trim_whitespace: bool) -> str: + if item is None: + text = "" + elif item_conversion == "keep strings as-is" and isinstance(item, str): + text = item + elif item_conversion == "convert dicts and lists to json text" and isinstance(item, (dict, list)): + text = _safe_json(item) + elif isinstance(item, str): + text = item + elif isinstance(item, (dict, list)): + text = _safe_json(item) + else: + text = str(item) + + if trim_whitespace and isinstance(text, str): + text = text.strip() + + return text + + @staticmethod + def _split_newline_items(raw_text: str) -> List[str]: + return raw_text.splitlines() + + @staticmethod + def _coerce_items(input_data: str, input_interpretation: str) -> tuple[List[Any], str]: + raw_text = "" if input_data is None else str(input_data) + stripped = raw_text.strip() + + if input_interpretation == "treat as plain text": + return [raw_text], "plain_text" + + if input_interpretation == "treat as newline list": + return ZMongoDataPassThroughNode._split_newline_items(raw_text), "newline_list" + + if input_interpretation == "treat as json": + if not stripped: + return [], "json_empty" + parsed = json.loads(stripped) + if isinstance(parsed, list): + return parsed, "json_list" + return [parsed], "json_object" + + # auto detect + if not stripped: + return [], "empty" + + try: + parsed = json.loads(stripped) + if isinstance(parsed, list): + return parsed, "json_list" + if isinstance(parsed, dict): + return [parsed], "json_object" + except Exception: + pass + + if "\n" in raw_text: + return ZMongoDataPassThroughNode._split_newline_items(raw_text), "newline_list" + + return [raw_text], "plain_text" + + def pass_through( + self, + input_data, + input_interpretation, + output_mode, + item_conversion, + skip_empty_items, + trim_whitespace, + ): + try: + raw_items, detected_input_type = self._coerce_items( + input_data=input_data, + input_interpretation=input_interpretation, + ) + + normalized_items: List[str] = [] + for item in raw_items: + text = self._normalize_item( + item=item, + item_conversion=item_conversion, + trim_whitespace=trim_whitespace, + ) + if skip_empty_items and not text: + continue + normalized_items.append(text) + + item_count = len(normalized_items) + normalized_json = _safe_json(normalized_items) + + if item_count == 0: + single_item = "" + list_items = [] + else: + single_item = normalized_items[0] + + if output_mode == "single item only": + list_items = [] + else: + # Important: return a flat list of strings here + list_items = normalized_items + + return ( + single_item, + list_items, + int(item_count), + detected_input_type, + normalized_json, + ) + + except Exception as exc: + error_text = f"ZMongoDataPassThroughNode error: {exc}" + return ( + error_text, + [], + 0, + "error", + _safe_json({"error": str(exc)}), + ) + + +class ZMongoSaveTextNode: + """ + Saves text or JSON text to file(s) in the ComfyUI output directory. + + User-oriented behavior: + - Save one value to one file + - Save a JSON list to separate files + - Save a true ComfyUI list to separate files + - Save many items into one combined file + - Overwrite exact filenames, append to files, or auto-create unique names + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "content_to_save": ("STRING", {"forceInput": True}), + "base_filename": ("STRING", {"default": "zmongo_output"}), + "starting_index": ("INT", {"default": 1, "min": 0, "max": 1000000}), + "file_type": (["txt", "json", "md"],), + "save_behavior": ([ + "write new content", + "append to existing file", + ],), + "save_layout": ([ + "one combined file", + "separate file per item", + ],), + "reuse_exact_filenames": ("BOOLEAN", {"default": False}), + } + } + + RETURN_TYPES = ("STRING", "INT") + RETURN_NAMES = ("saved_file_path_or_paths", "saved_file_count") + FUNCTION = "save_text" + CATEGORY = "ZMongo/Output" + INPUT_IS_LIST = True + + @staticmethod + def _unwrap_scalar(value, default=None): + """ + ComfyUI may pass scalar inputs as one-item lists when INPUT_IS_LIST = True. + """ + if isinstance(value, list): + if not value: + return default + return value[0] + return value + + @classmethod + def _unwrap_and_cast_int(cls, value, default=0) -> int: + value = cls._unwrap_scalar(value, default) + if value is None: + return int(default) + return int(value) + + @classmethod + def _unwrap_and_cast_bool(cls, value, default=False) -> bool: + value = cls._unwrap_scalar(value, default) + if isinstance(value, bool): + return value + return bool(value) + + @classmethod + def _unwrap_and_cast_str(cls, value, default="") -> str: + value = cls._unwrap_scalar(value, default) + if value is None: + return str(default) + return str(value) + + @staticmethod + def _coerce_to_items(content_to_save): + """ + Accepts: + - plain string + - JSON list string + - real ComfyUI list input + """ + if isinstance(content_to_save, list): + if len(content_to_save) == 1 and isinstance(content_to_save[0], str): + raw = content_to_save[0] + stripped = raw.strip() + if stripped: + try: + parsed = json.loads(stripped) + if isinstance(parsed, list): + return parsed + except Exception: + pass + return content_to_save + + if content_to_save is None: + return [""] + + raw = str(content_to_save) + stripped = raw.strip() + if not stripped: + return [""] + + try: + parsed = json.loads(stripped) + if isinstance(parsed, list): + return parsed + return [raw] + except Exception: + return [raw] + + @staticmethod + def _item_to_text(item, file_type: str) -> str: + if item is None: + return "" + if isinstance(item, str): + return item + if file_type == "json": + return json.dumps(item, indent=2, ensure_ascii=False) + if isinstance(item, (dict, list)): + return json.dumps(item, indent=2, ensure_ascii=False) + return str(item) + + @staticmethod + def _ensure_parent_dir(path: str): + import os + parent = os.path.dirname(path) + if parent: + os.makedirs(parent, exist_ok=True) + + @staticmethod + def _make_single_filename(base_filename: str, file_type: str) -> str: + return f"{base_filename}.{file_type}" + + @staticmethod + def _make_indexed_filename(base_filename: str, index: int, file_type: str) -> str: + return f"{base_filename}_{int(index):04d}.{file_type}" + + @staticmethod + def _make_unique_path(path: str, file_type: str) -> str: + import os + + if not os.path.exists(path): + return path + + base = path[: -(len(file_type) + 1)] + counter = 1 + while True: + candidate = f"{base}_{counter}.{file_type}" + if not os.path.exists(candidate): + return candidate + counter += 1 + + def save_text( + self, + content_to_save, + base_filename, + starting_index, + file_type, + save_behavior, + save_layout, + reuse_exact_filenames, + ): + try: + import os + import folder_paths + + # Normalize scalar inputs that may arrive as [value] + base_filename = self._unwrap_and_cast_str(base_filename, "zmongo_output") + starting_index = self._unwrap_and_cast_int(starting_index, 1) + file_type = self._unwrap_and_cast_str(file_type, "txt") + save_behavior = self._unwrap_and_cast_str(save_behavior, "write new content") + save_layout = self._unwrap_and_cast_str(save_layout, "one combined file") + reuse_exact_filenames = self._unwrap_and_cast_bool(reuse_exact_filenames, False) + + output_dir = folder_paths.get_output_directory() + items = self._coerce_to_items(content_to_save) + + append_mode = (save_behavior == "append to existing file") + file_open_mode = "a" if append_mode else "w" + + saved_paths = [] + + if save_layout == "one combined file": + filename = self._make_single_filename(base_filename, file_type) + full_path = os.path.join(output_dir, filename) + + if not reuse_exact_filenames and not append_mode: + full_path = self._make_unique_path(full_path, file_type) + + self._ensure_parent_dir(full_path) + + with open(full_path, file_open_mode, encoding="utf-8") as f: + for i, item in enumerate(items): + text_value = self._item_to_text(item, file_type) + f.write(text_value) + if i < len(items) - 1: + f.write("\n") + + saved_paths.append(full_path) + + else: + for offset, item in enumerate(items): + item_index = int(starting_index) + offset + filename = self._make_indexed_filename(base_filename, item_index, file_type) + full_path = os.path.join(output_dir, filename) + + if not reuse_exact_filenames and not append_mode: + full_path = self._make_unique_path(full_path, file_type) + + self._ensure_parent_dir(full_path) + + text_value = self._item_to_text(item, file_type) + + with open(full_path, file_open_mode, encoding="utf-8") as f: + f.write(text_value) + + saved_paths.append(full_path) + + if len(saved_paths) == 1: + return (saved_paths[0], len(saved_paths)) + + return (json.dumps(saved_paths, indent=2, ensure_ascii=False), len(saved_paths)) + + except Exception as exc: + return (f"Error saving file(s): {exc}", 0) + + +class ZMongoPromptDemoNode(ZMongoAPIMixin): + """ + Standalone demo node for real-database prompt retrieval. + + Purpose: + - Connect to a real MongoDB collection + - Fetch one record by 1-based index using a stable sort + - Extract prompt text from a selected field path with fallbacks + - Return prompt text plus useful demo metadata + + Intended use: + - prompt_text -> CLIPTextEncode (positive) + - summary / selected_record_json -> PreviewAny + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "mongo_uri": ("STRING", {"default": "mongodb://127.0.0.1:27017"}), + "database_name": ("STRING", {"default": "test"}), + "collection_name": ("STRING", {"default": "ocr_docs"}), + "select_index": ("INT", {"default": 1, "min": 1, "max": 100000000}), + "sort_field": ("STRING", {"default": "_id"}), + "sort_direction": (["ascending", "descending"],), + "field_path": ("STRING", {"default": "text"}), + "fallback_field_paths": ("STRING", { + "default": "text\ndoc_text\ndocument.text\nprompt", + "multiline": True, + }), + "filter_json": ("STRING", {"default": "{}", "multiline": True}), + "projection_json": ("STRING", {"default": "{}", "multiline": True}), + "strip_prompt": ("BOOLEAN", {"default": True}), + "max_prompt_length": ("INT", {"default": 2000, "min": 0, "max": 100000}), + "refresh_nonce": ("INT", {"default": 0, "min": 0, "max": 999999999}), + } + } + + RETURN_TYPES = ( + "STRING", # prompt_text + "STRING", # selected_record_json + "STRING", # summary + "INT", # selected_index_out + "INT", # total_count + "STRING", # record_id + ) + RETURN_NAMES = ( + "prompt_text", + "selected_record_json", + "summary", + "selected_index_out", + "total_count", + "record_id", + ) + FUNCTION = "fetch_demo_prompt" + CATEGORY = "ZMongo/Demo" + + @staticmethod + def _parse_json_object(raw: str, field_name: str) -> Dict[str, Any]: + text = (raw or "").strip() + if not text: + return {} + try: + parsed = json.loads(text) + except Exception as exc: + raise ValueError(f"{field_name} is not valid JSON: {exc}") from exc + if not isinstance(parsed, dict): + raise ValueError(f"{field_name} must be a JSON object") + return parsed + + @staticmethod + def _normalize_sort(sort_field: str, sort_direction: str): + direction = 1 if sort_direction == "ascending" else -1 + field = (sort_field or "").strip() or "_id" + return [(field, direction)] + + @staticmethod + def _candidate_paths(primary: str, fallback_blob: str) -> List[str]: + paths: List[str] = [] + if primary and primary.strip(): + paths.append(primary.strip()) + for line in (fallback_blob or "").splitlines(): + value = line.strip() + if value and value not in paths: + paths.append(value) + return paths + + @staticmethod + def _extract_prompt_from_record( + record: Dict[str, Any], + paths: List[str], + strip_prompt: bool, + max_prompt_length: int, + ) -> (str, str): + for path in paths: + try: + value = DataProcessor.get_value(record, path) + except Exception: + value = None + + if value is None: + continue + + if isinstance(value, (dict, list)): + text = _safe_json(value) + else: + text = str(value) + + if strip_prompt: + text = text.strip() + + if max_prompt_length > 0 and len(text) > max_prompt_length: + text = text[:max_prompt_length] + + return text, path + + return "", "" + + @staticmethod + def _record_id(record: Dict[str, Any]) -> str: + value = record.get("_id") + return "" if value is None else str(value) + + @staticmethod + def _short_preview(text: str, length: int = 180) -> str: + preview = (text or "").replace("\n", " ").strip() + if len(preview) > length: + return preview[:length] + "..." + return preview + + def fetch_demo_prompt( + self, + mongo_uri, + database_name, + collection_name, + select_index, + sort_field, + sort_direction, + field_path, + fallback_field_paths, + filter_json, + projection_json, + strip_prompt, + max_prompt_length, + refresh_nonce, + ): + _ = refresh_nonce + client = None + try: + client = MongoClient(mongo_uri, serverSelectionTimeoutMS=5000) + client.admin.command("ping") + + collection = client[database_name][collection_name] + mongo_filter = self._parse_json_object(filter_json, "filter_json") + projection = self._parse_json_object(projection_json, "projection_json") + sort_spec = self._normalize_sort(sort_field, sort_direction) + + total_count = collection.count_documents(mongo_filter) + if total_count == 0: + empty_payload = { + "status": "empty", + "database_name": database_name, + "collection_name": collection_name, + "filter_json": mongo_filter, + } + return ( + "", + json.dumps(empty_payload, indent=2), + f"No records found in {database_name}.{collection_name}", + 0, + 0, + "", + ) + + bounded_index = max(1, min(int(select_index), int(total_count))) + + cursor = collection.find(mongo_filter, projection or None).sort(sort_spec).skip(bounded_index - 1).limit(1) + selected_record = next(cursor, None) + + if selected_record is None: + error_payload = { + "error": "Selected record could not be retrieved", + "database_name": database_name, + "collection_name": collection_name, + "selected_index": bounded_index, + } + return ( + "", + json.dumps(error_payload, indent=2), + "Selected record could not be retrieved.", + int(bounded_index), + int(total_count), + "", + ) + + paths = self._candidate_paths(field_path, fallback_field_paths) + prompt_text, used_path = self._extract_prompt_from_record( + record=selected_record, + paths=paths, + strip_prompt=strip_prompt, + max_prompt_length=max_prompt_length, + ) + + record_copy = dict(selected_record) + record_copy["_selected_index"] = int(bounded_index) + record_copy["_database_name"] = str(database_name) + record_copy["_collection_name"] = str(collection_name) + record_copy["_field_path_used"] = used_path + + record_id = self._record_id(record_copy) + + if not prompt_text: + summary = ( + f"Demo prompt fetch succeeded, but no prompt text was found.\n" + f"Collection: {database_name}.{collection_name}\n" + f"Selected index: {bounded_index} / {total_count}\n" + f"Record _id: {record_id}\n" + f"Tried paths: {', '.join(paths) if paths else '(none)'}" + ) + return ( + "", + _safe_json(record_copy), + summary, + int(bounded_index), + int(total_count), + record_id, + ) + + summary = ( + f"Real DB prompt demo\n" + f"Collection: {database_name}.{collection_name}\n" + f"Selected index: {bounded_index} / {total_count}\n" + f"Record _id: {record_id}\n" + f"Prompt path used: {used_path or '(none)'}\n" + f"Prompt preview: {self._short_preview(prompt_text)}" + ) + + return ( + prompt_text, + _safe_json(record_copy), + summary, + int(bounded_index), + int(total_count), + record_id, + ) + + except Exception as exc: + error_payload = { + "error": str(exc), + "mongo_uri": mongo_uri, + "database_name": database_name, + "collection_name": collection_name, + } + return ( + "", + json.dumps(error_payload, indent=2), + f"ZMongoPromptDemoNode error: {exc}", + 0, + 0, + "", + ) + finally: + if client is not None: + try: + client.close() + except Exception: + pass + + +class ZMongoLoopControllerNode: + """ + Pure loop/index controller node. + + This node does not touch MongoDB. + It only computes which record index should be used on this pass + and exposes a loop counter plus status flags. + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "start_index": ("INT", {"default": 1, "min": 1, "max": 100000000}), + "end_index": ("INT", {"default": 10, "min": 1, "max": 100000000}), + "step": ("INT", {"default": 1, "min": 1, "max": 100000}), + "loop_count": ("INT", {"default": 0, "min": 0, "max": 100000000}), + "loop_mode": (["stop_at_end", "wrap"],), + "reset_nonce": ("INT", {"default": 0, "min": 0, "max": 999999999}), + "max_iterations": ("INT", {"default": 1000, "min": 1, "max": 100000000}), + } + } + + RETURN_TYPES = ( + "INT", # record_index + "INT", # next_loop_count + "BOOLEAN", # has_next + "BOOLEAN", # is_finished + "STRING", # summary + "STRING", # loop_state_json + ) + RETURN_NAMES = ( + "record_index", + "next_loop_count", + "has_next", + "is_finished", + "summary", + "loop_state_json", + ) + FUNCTION = "control_loop" + CATEGORY = "ZMongo/Looping" + + def control_loop( + self, + start_index, + end_index, + step, + loop_count, + loop_mode, + reset_nonce, + max_iterations, + ): + _ = reset_nonce + + safe_start = max(1, int(start_index)) + safe_end = max(1, int(end_index)) + safe_step = max(1, int(step)) + safe_loop_count = max(0, int(loop_count)) + safe_max_iterations = max(1, int(max_iterations)) + + forward = safe_start <= safe_end + + if forward: + sequence = list(range(safe_start, safe_end + 1, safe_step)) + else: + sequence = list(range(safe_start, safe_end - 1, -safe_step)) + + if not sequence: + state = { + "start_index": safe_start, + "end_index": safe_end, + "step": safe_step, + "loop_count": safe_loop_count, + "record_index": 0, + "has_next": False, + "is_finished": True, + "loop_mode": loop_mode, + "max_iterations": safe_max_iterations, + "sequence_length": 0, + } + return ( + 0, + safe_loop_count, + False, + True, + "Loop controller: empty sequence.", + _safe_json(state), + ) + + seq_len = len(sequence) + + if safe_loop_count >= safe_max_iterations: + record_index = sequence[-1] + state = { + "start_index": safe_start, + "end_index": safe_end, + "step": safe_step, + "loop_count": safe_loop_count, + "record_index": record_index, + "has_next": False, + "is_finished": True, + "loop_mode": loop_mode, + "max_iterations": safe_max_iterations, + "sequence_length": seq_len, + } + return ( + int(record_index), + int(safe_loop_count), + False, + True, + f"Loop controller reached max_iterations={safe_max_iterations}.", + _safe_json(state), + ) + + if loop_mode == "wrap": + seq_pos = safe_loop_count % seq_len + record_index = sequence[seq_pos] + has_next = True + is_finished = False + next_loop_count = safe_loop_count + 1 + else: + seq_pos = min(safe_loop_count, seq_len - 1) + record_index = sequence[seq_pos] + has_next = safe_loop_count < (seq_len - 1) + is_finished = not has_next + next_loop_count = safe_loop_count + 1 if has_next else safe_loop_count + + state = { + "start_index": safe_start, + "end_index": safe_end, + "step": safe_step, + "loop_count": safe_loop_count, + "record_index": int(record_index), + "sequence_position": int(seq_pos), + "sequence_length": int(seq_len), + "has_next": bool(has_next), + "is_finished": bool(is_finished), + "loop_mode": str(loop_mode), + "max_iterations": int(safe_max_iterations), + } + + summary = ( + f"Loop count: {safe_loop_count}\n" + f"Record index: {record_index}\n" + f"Range: {safe_start} -> {safe_end} (step {safe_step})\n" + f"Mode: {loop_mode}\n" + f"Has next: {has_next}\n" + f"Finished: {is_finished}" + ) + + return ( + int(record_index), + int(next_loop_count), + bool(has_next), + bool(is_finished), + summary, + _safe_json(state), + ) + + +class ZMongoRecordLoopNode: + """ + User-friendly record selection node for ZMongo. + + Main purposes: + - Get one record + - Get a range of records + - Get all records + - Get a batch from a range + - Get random records + - Return one field value, one record, or JSON lists for saving/export + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "mongo_uri": ("STRING", {"default": "mongodb://127.0.0.1:27017"}), + "database_name": ("STRING", {"default": "test"}), + "collection_name": ("STRING", {"default": "ocr_docs"}), + + "selection_mode": ([ + "one record", + "record range", + "all records", + "first batch from range", + "one random record", + "random batch", + ],), + + "record_start": ("INT", {"default": 1, "min": 1, "max": 100000000}), + "record_end": ("INT", {"default": 10, "min": 1, "max": 100000000}), + "record_step": ("INT", {"default": 1, "min": 1, "max": 100000}), + "records_per_batch": ("INT", {"default": 1, "min": 1, "max": 1024}), + + "sort_by_field": ("STRING", {"default": "_id"}), + "sort_order": (["ascending", "descending"],), + + "primary_field_name": ("STRING", {"default": "text"}), + "fallback_field_names": ("STRING", { + "default": "text\ndoc_text\ndocument.text\nprompt", + "multiline": True, + }), + + "filter_query_json": ("STRING", {"default": "{}", "multiline": True}), + "fields_to_include_json": ("STRING", {"default": "{}", "multiline": True}), + + "randomize_order": ("BOOLEAN", {"default": False}), + "random_seed": ("INT", {"default": 0, "min": 0, "max": 2147483647}), + "remove_duplicate_ids": ("BOOLEAN", {"default": True}), + "skip_records_with_empty_field": ("BOOLEAN", {"default": True}), + "trim_whitespace": ("BOOLEAN", {"default": True}), + "maximum_field_length": ("INT", {"default": 2000, "min": 0, "max": 100000}), + + "empty_value_behavior": ([ + "skip record", + "return empty string", + "use placeholder text", + ],), + "placeholder_text": ("STRING", {"default": "empty prompt"}), + + "range_end_behavior": ([ + "stop at end", + "wrap to beginning", + "bounce back", + ],), + + "output_format": ([ + "one field value", + "one record as json", + "one record and field value", + "record list as json", + "field value list as json", + ],), + + "refresh_nonce": ("INT", {"default": 0, "min": 0, "max": 999999999}), + } + } + + RETURN_TYPES = ( + "STRING", # selected_record_json_or_list + "STRING", # selected_field_value_or_list + "INT", # current_record_index + "INT", # total_record_count + "INT", # returned_item_count + "STRING", # summary + "STRING", # selection_state_json + ) + RETURN_NAMES = ( + "selected_record_json_or_list", + "selected_field_value_or_list", + "current_record_index", + "total_record_count", + "returned_item_count", + "summary", + "selection_state_json", + ) + FUNCTION = "iterate_records" + CATEGORY = "ZMongo/Looping" + + @staticmethod + def _parse_json_object(raw: str, field_name: str) -> Dict[str, Any]: + text = (raw or "").strip() + if not text: + return {} + try: + parsed = json.loads(text) + except Exception as exc: + raise ValueError(f"{field_name} is not valid JSON: {exc}") from exc + if not isinstance(parsed, dict): + raise ValueError(f"{field_name} must be a JSON object") + return parsed + + @staticmethod + def _fallback_paths(primary_field_name: str, fallback_field_names: str) -> List[str]: + paths: List[str] = [] + if primary_field_name and primary_field_name.strip(): + paths.append(primary_field_name.strip()) + for line in (fallback_field_names or "").splitlines(): + value = line.strip() + if value and value not in paths: + paths.append(value) + return paths + + @staticmethod + def _normalize_sort(sort_by_field: str, sort_order: str): + direction = 1 if sort_order == "ascending" else -1 + field = (sort_by_field or "").strip() or "_id" + return [(field, direction)] + + @staticmethod + def _extract_field_value( + record: Dict[str, Any], + paths: List[str], + trim_whitespace: bool, + maximum_field_length: int, + ) -> tuple[str, str]: + for path in paths: + try: + value = DataProcessor.get_value(record, path) + except Exception: + value = None + + if value is None: + continue + + if isinstance(value, (dict, list)): + try: + text = _safe_json(value) + except Exception: + text = str(value) + else: + text = str(value) + + if trim_whitespace: + text = text.strip() + + if maximum_field_length > 0 and len(text) > maximum_field_length: + text = text[:maximum_field_length] + + return text, path + + return "", "" + + @staticmethod + def _apply_empty_value_behavior( + field_value: str, + empty_value_behavior: str, + placeholder_text: str, + ) -> str: + if field_value: + return field_value + if empty_value_behavior == "return empty string": + return "" + if empty_value_behavior == "use placeholder text": + return placeholder_text or "empty prompt" + return "" + + @staticmethod + def _doc_id_string(doc: Dict[str, Any]) -> str: + value = doc.get("_id") + return "" if value is None else str(value) + + def _load_records( + self, + collection, + filter_query: Dict[str, Any], + fields_to_include: Dict[str, Any], + sort_spec, + ) -> List[Dict[str, Any]]: + cursor = collection.find(filter_query, fields_to_include or None).sort(sort_spec) + return list(cursor) + + def _select_records( + self, + records: List[Dict[str, Any]], + selection_mode: str, + record_start: int, + record_end: int, + record_step: int, + records_per_batch: int, + randomize_order: bool, + random_seed: int, + remove_duplicate_ids: bool, + ) -> tuple[List[Dict[str, Any]], int]: + if not records: + return [], 0 + + working = list(records) + + if remove_duplicate_ids: + seen = set() + deduped = [] + for record in working: + sid = self._doc_id_string(record) + key = sid or id(record) + if key in seen: + continue + seen.add(key) + deduped.append(record) + working = deduped + + rng = random.Random(random_seed) + + if randomize_order: + rng.shuffle(working) + + total = len(working) + + bounded_start = max(1, min(int(record_start), total)) + bounded_end = max(1, min(int(record_end), total)) + safe_step = max(1, int(record_step)) + safe_batch_size = max(1, int(records_per_batch)) + + if selection_mode == "one record": + selected = [working[bounded_start - 1]] + return selected, bounded_start + + if selection_mode == "one random record": + idx = rng.randint(1, total) + selected = [working[idx - 1]] + return selected, idx + + if selection_mode == "all records": + return working, 1 if working else 0 + + if selection_mode == "record range": + if bounded_start <= bounded_end: + indices = list(range(bounded_start, bounded_end + 1, safe_step)) + else: + indices = list(range(bounded_start, bounded_end - 1, -safe_step)) + selected = [working[i - 1] for i in indices] + return selected, indices[0] if indices else 0 + + if selection_mode == "first batch from range": + if bounded_start <= bounded_end: + indices = list(range(bounded_start, bounded_end + 1, safe_step)) + else: + indices = list(range(bounded_start, bounded_end - 1, -safe_step)) + indices = indices[:safe_batch_size] + selected = [working[i - 1] for i in indices] + return selected, indices[0] if indices else 0 + + if selection_mode == "random batch": + population = list(range(1, total + 1)) + rng.shuffle(population) + indices = sorted(population[:safe_batch_size]) + selected = [working[i - 1] for i in indices] + return selected, indices[0] if indices else 0 + + return [], 0 + + def iterate_records( + self, + mongo_uri, + database_name, + collection_name, + selection_mode, + record_start, + record_end, + record_step, + records_per_batch, + sort_by_field, + sort_order, + primary_field_name, + fallback_field_names, + filter_query_json, + fields_to_include_json, + randomize_order, + random_seed, + remove_duplicate_ids, + skip_records_with_empty_field, + trim_whitespace, + maximum_field_length, + empty_value_behavior, + placeholder_text, + range_end_behavior, + output_format, + refresh_nonce, + ): + _ = refresh_nonce + client = None + try: + client = MongoClient(mongo_uri, serverSelectionTimeoutMS=5000) + client.admin.command("ping") + collection = client[database_name][collection_name] + + total_record_count = collection.count_documents({}) + + filter_query = self._parse_json_object(filter_query_json, "filter_query_json") + fields_to_include = self._parse_json_object(fields_to_include_json, "fields_to_include_json") + sort_spec = self._normalize_sort(sort_by_field, sort_order) + paths = self._fallback_paths(primary_field_name, fallback_field_names) + + records = self._load_records(collection, filter_query, fields_to_include, sort_spec) + + selected_records, current_record_index = self._select_records( + records=records, + selection_mode=selection_mode, + record_start=record_start, + record_end=record_end, + record_step=record_step, + records_per_batch=records_per_batch, + randomize_order=randomize_order, + random_seed=random_seed, + remove_duplicate_ids=remove_duplicate_ids, + ) + + if not selected_records: + selection_state = { + "database_name": str(database_name), + "collection_name": str(collection_name), + "total_record_count": int(total_record_count), + "filtered_record_count": int(len(records)), + "record_start": int(record_start), + "record_end": int(record_end), + "record_step": int(record_step), + "current_record_index": 0, + "records_per_batch": int(records_per_batch), + "selection_mode": str(selection_mode), + "range_end_behavior": str(range_end_behavior), + "sort_by_field": str(sort_by_field or "_id"), + "sort_order": str(sort_order), + "field_name_used": "", + "record_id": "", + "has_next": False, + "has_previous": False, + "wrapped": False, + "direction": "forward", + } + return ( + "[]", + "", + 0, + int(total_record_count), + 0, + f"No records matched in {database_name}.{collection_name}", + _safe_json(selection_state), + ) + + field_items: List[Dict[str, Any]] = [] + kept_records: List[Dict[str, Any]] = [] + + for offset, record in enumerate(selected_records): + field_value, used_path = self._extract_field_value( + record=record, + paths=paths, + trim_whitespace=trim_whitespace, + maximum_field_length=maximum_field_length, + ) + field_value = self._apply_empty_value_behavior( + field_value=field_value, + empty_value_behavior=empty_value_behavior, + placeholder_text=placeholder_text, + ) + + if skip_records_with_empty_field and not field_value: + continue + + record_copy = dict(record) + record_copy["_loop_index"] = int(current_record_index + offset) + record_copy["_database_name"] = str(database_name) + record_copy["_collection_name"] = str(collection_name) + record_copy["_field_path_used"] = used_path + + kept_records.append(record_copy) + field_items.append({ + "index": int(current_record_index + offset), + "record_id": self._doc_id_string(record_copy), + "field_path_used": used_path, + "field_value": field_value, + }) + + if not kept_records: + selection_state = { + "database_name": str(database_name), + "collection_name": str(collection_name), + "total_record_count": int(total_record_count), + "filtered_record_count": int(len(records)), + "record_start": int(record_start), + "record_end": int(record_end), + "record_step": int(record_step), + "current_record_index": int(current_record_index), + "records_per_batch": int(records_per_batch), + "selection_mode": str(selection_mode), + "range_end_behavior": str(range_end_behavior), + "sort_by_field": str(sort_by_field or "_id"), + "sort_order": str(sort_order), + "field_name_used": "", + "record_id": "", + "has_next": False, + "has_previous": current_record_index > 1, + "wrapped": False, + "direction": "forward", + } + return ( + "[]", + "", + int(current_record_index), + int(total_record_count), + 0, + "All selected records were skipped because the chosen field was empty.", + _safe_json(selection_state), + ) + + first_record = kept_records[0] + first_field_value = field_items[0]["field_value"] + first_path = field_items[0]["field_path_used"] + filtered_record_count = len(records) + + selection_state = { + "database_name": str(database_name), + "collection_name": str(collection_name), + "total_record_count": int(total_record_count), + "filtered_record_count": int(filtered_record_count), + "record_start": int(record_start), + "record_end": int(record_end), + "record_step": int(record_step), + "current_record_index": int(field_items[0]["index"]), + "current_zero_based": max(0, int(field_items[0]["index"]) - 1), + "records_per_batch": int(records_per_batch), + "batch_number": 1, + "has_next": int(field_items[-1]["index"]) < int(filtered_record_count), + "has_previous": int(field_items[0]["index"]) > 1, + "wrapped": False, + "direction": "forward", + "sort_by_field": str(sort_by_field or "_id"), + "sort_order": str(sort_order), + "selection_mode": str(selection_mode), + "range_end_behavior": str(range_end_behavior), + "field_name_used": first_path, + "record_id": self._doc_id_string(first_record), + } + + field_preview = first_field_value.replace("\n", " ") + if len(field_preview) > 180: + field_preview = field_preview[:180] + "..." + + summary = ( + f"Selection mode: {selection_mode}\n" + f"Collection: {database_name}.{collection_name}\n" + f"Filtered records: {filtered_record_count} / {total_record_count}\n" + f"Current record index: {field_items[0]['index']}\n" + f"Record _id: {self._doc_id_string(first_record)}\n" + f"Field used: {first_path or '(none)'}\n" + f"Value preview: {field_preview}" + ) + + if output_format == "one record as json": + selected_record_output = _safe_json(first_record) + selected_field_output = first_field_value + returned_item_count = 1 + elif output_format == "one field value": + selected_record_output = _safe_json(first_record) + selected_field_output = first_field_value + returned_item_count = 1 + elif output_format == "one record and field value": + selected_record_output = _safe_json(first_record) + selected_field_output = first_field_value + returned_item_count = 1 + elif output_format == "record list as json": + selected_record_output = _safe_json(kept_records) + selected_field_output = _safe_json(field_items) + returned_item_count = len(kept_records) + elif output_format == "field value list as json": + selected_record_output = _safe_json(kept_records) + selected_field_output = _safe_json([item["field_value"] for item in field_items]) + returned_item_count = len(field_items) + else: + selected_record_output = _safe_json(first_record) + selected_field_output = first_field_value + returned_item_count = 1 + + return ( + selected_record_output, + selected_field_output, + int(field_items[0]["index"]), + int(total_record_count), + int(returned_item_count), + summary, + _safe_json(selection_state), + ) + + except Exception as exc: + error_payload = { + "error": str(exc), + "mongo_uri": mongo_uri, + "database_name": database_name, + "collection_name": collection_name, + } + return ( + json.dumps(error_payload, indent=2), + "", + 0, + 0, + 0, + f"ZMongoRecordLoopNode error: {exc}", + json.dumps(error_payload, indent=2), + ) + finally: + if client is not None: + try: + client.close() + except Exception: + pass + + +class ZRetrieverNode: + """Semantic search node that avoids running ad-hoc event loops inside ComfyUI.""" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "query_text": ("STRING", {"multiline": True}), + "collection": ("STRING", {"default": "ocr_docs"}), + "similarity_threshold": ("FLOAT", {"default": 0.7, "min": 0.0, "max": 1.0}), + "n_results": ("INT", {"default": 3, "min": 1, "max": 50}), + } + } + + RETURN_TYPES = ("STRING", "STRING") + RETURN_NAMES = ("top_result_json", "summary") + FUNCTION = "retrieve" + CATEGORY = "ZMongo/AI" + + def retrieve(self, query_text, collection, similarity_threshold, n_results): + # 1. Guard against empty/whitespace queries + if not query_text or not query_text.strip(): + return ("{}", "Search skipped: Query text is empty.") + + if ZEmbedder is None: + return ("{}", "Search failed: ZEmbedder is not importable.") + + try: + zmongo = _get_zmongo() + # Ensure dimensionality matches your BGE-M3 config + embedder = ZEmbedder(output_dimensionality=768) + + result = zmongo.run_sync( + embedder.find_similar_documents, + query_text=query_text, + target_collection=collection, + n_results=n_results, + ) + if not isinstance(result, SafeResult): + result = SafeResult.fail(f"Unexpected search result type: {type(result).__name__}") + + if not result.success: + return ("{}", f"Search failed: {result.error}") + + raw_hits = (result.data or {}).get("results", []) + hits = [ + hit for hit in raw_hits + if float(hit.get("retrieval_score", 0.0)) >= float(similarity_threshold) + ] + + if not hits: + return ("{}", "No matching documents found above the similarity threshold.") + + return _safe_json(hits[0].get("document", {})), _build_summary_from_hits(hits) + except Exception as exc: + logger.exception("ZRetrieverNode failure") + return ("{}", f"Search failed: {exc}") + + +class ZMongoDatabaseBrowserNode: + """Synchronous browser node for listing records from a collection.""" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "mongo_uri": ("STRING", {"default": "mongodb://127.0.0.1:27017", "multiline": False}), + "database_name": ("STRING", {"default": "test", "multiline": False}), + "collection_name": ("STRING", {"default": "ocr_docs", "multiline": False}), + "limit": ("INT", {"default": 50, "min": 1, "max": 1000, "step": 1}), + "select_index": ("INT", {"default": 1, "min": 1, "max": 100000, "step": 1}), + "refresh_nonce": ("INT", {"default": 1, "min": 0, "max": 999999999, "step": 1}), + } + } + + RETURN_TYPES = ("STRING", "STRING", "INT", "STRING", "STRING", "INT") + RETURN_NAMES = ( + "selected_record_json", + "record_list_summary", + "total_count", + "db_name_out", + "coll_name_out", + "selected_index_out", + ) + FUNCTION = "browse_database" + CATEGORY = "ZMongo/Database" + + @staticmethod + def _make_summary(records: List[dict]) -> str: + if not records: + return "No records found." + + lines = [] + for idx, record in enumerate(records, start=1): + record_id = record.get("_id", "") + title = ( + record.get("title") + or record.get("name") + or record.get("username") + or record.get("email") + or record.get("text") + or record.get("content") + or "" + ) + if isinstance(title, str): + title = title.strip().replace("\n", " ") + if len(title) > 100: + title = title[:100] + "..." + else: + title = str(title) + if not title: + title = f"keys={list(record.keys())[:6]}" + lines.append(f"{idx}. _id={record_id} | {title}") + return "\n".join(lines) + + def browse_database(self, mongo_uri, database_name, collection_name, limit, select_index, refresh_nonce): + _ = refresh_nonce + client = None + try: + client = MongoClient(mongo_uri, serverSelectionTimeoutMS=5000) + client.admin.command("ping") + collection = client[database_name][collection_name] + + total_count = collection.count_documents({}) + + if total_count == 0: + return ( + json.dumps( + { + "status": "empty", + "message": f"No records found in {database_name}.{collection_name}", + }, + indent=2, + ), + "No records found.", + 0, + str(database_name), + str(collection_name), + 0, + ) + + # Clamp the requested index to a valid 1-based record index. + bounded_index = max(1, min(int(select_index), int(total_count))) + + # Use a stable sort so index N always maps to the same record ordering. + sort_spec = [("_id", 1)] + + # Fetch the selected record directly by index. + selected_cursor = collection.find({}).sort(sort_spec).skip(bounded_index - 1).limit(1) + selected_record = next(selected_cursor, None) + + if selected_record is None: + return ( + json.dumps( + { + "error": "Selected record could not be retrieved", + "selected_index": bounded_index, + }, + indent=2, + ), + "Database browse error: Selected record could not be retrieved.", + int(total_count), + str(database_name), + str(collection_name), + int(bounded_index), + ) + + # Fetch a preview list for the browser summary. + summary_limit = max(1, min(int(limit), int(total_count))) + records = list(collection.find({}).sort(sort_spec).limit(summary_limit)) + + summary_lines = [] + for idx, record in enumerate(records, start=1): + record_id = record.get("_id", "") + title = ( + record.get("title") + or record.get("name") + or record.get("username") + or record.get("email") + or record.get("text") + or record.get("content") + or "" + ) + if isinstance(title, str): + title = title.strip().replace("\n", " ") + if len(title) > 100: + title = title[:100] + "..." + else: + title = str(title) + + if not title: + title = f"keys={list(record.keys())[:6]}" + + marker = ">> " if idx == bounded_index else " " + summary_lines.append(f"{marker}{idx}. _id={record_id} | {title}") + + record_list_summary = "\n".join(summary_lines) if summary_lines else "No records found." + + # Include the selected index in the returned payload so downstream nodes + # can verify they are using the correct selected record. + selected_payload = dict(selected_record) + selected_payload["_selected_index"] = bounded_index + selected_payload["_database_name"] = str(database_name) + selected_payload["_collection_name"] = str(collection_name) + + return ( + _safe_json(selected_payload), + record_list_summary, + int(total_count), + str(database_name), + str(collection_name), + int(bounded_index), + ) + + except Exception as exc: + error_payload = { + "error": str(exc), + "mongo_uri": mongo_uri, + "database_name": database_name, + "collection_name": collection_name, + } + return ( + json.dumps(error_payload, indent=2), + f"Database browse error: {exc}", + 0, + str(database_name), + str(collection_name), + int(select_index), + ) + finally: + if client is not None: + try: + client.close() + except Exception: + pass + +class ZMongoRecordSplitter: + """Converts JSON arrays into a list output usable by downstream ComfyUI nodes.""" + + @classmethod + def INPUT_TYPES(cls): + return {"required": {"records_json": ("STRING", {"multiline": True})}} + + RETURN_TYPES = ("STRING", "INT") + RETURN_NAMES = ("record_json_list", "count") + OUTPUT_IS_LIST = (True, False) + FUNCTION = "split" + CATEGORY = "ZMongo/Database" + + def split(self, records_json): + parsed = _coerce_records(records_json) + if not parsed.success: + return ([json.dumps(parsed.to_dict())], 0) + + output = [DataProcessor.to_json(record) for record in parsed.data or []] + return (output, len(output)) + + +class ZMongoFieldSelector: + """ + Dot-path field extractor with frontend-assisted dropdown population. + + How it works: + - record_json comes from an upstream node + - on execution, this node flattens the JSON into path-like keys + - it returns those paths both as an output string and as UI metadata + - frontend JS silently updates the field_path dropdown choices + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "record_json": ("STRING", {"forceInput": True}), + "field_path": (["text"], {"default": "text"}), + } + } + + RETURN_TYPES = ("STRING", "STRING", "STRING") + RETURN_NAMES = ("field_value", "available_paths", "selected_path") + FUNCTION = "select_field" + CATEGORY = "ZMongo/Database" + + @staticmethod + def _normalize_available_paths(record: Dict[str, Any]) -> List[str]: + try: + flattened = DataProcessor.flatten_dict(record) + if isinstance(flattened, dict) and flattened: + paths = sorted(str(k) for k in flattened.keys()) + else: + paths = [] + except Exception: + paths = [] + + if not paths and isinstance(record, dict): + paths = sorted(str(k) for k in record.keys()) + + if not paths: + paths = ["text"] + + return paths + + @staticmethod + def _serialize_value(value: Any) -> str: + if value is None: + return "" + if isinstance(value, (dict, list)): + return _safe_json(value) + return str(value) + + def select_field(self, record_json, field_path): + try: + record = json.loads(record_json) + if not isinstance(record, dict): + raise ValueError("record_json must decode to a JSON object") + + available_paths_list = self._normalize_available_paths(record) + available_paths_text = "\n".join(available_paths_list) + + selected_path = field_path if field_path in available_paths_list else available_paths_list[0] + + try: + value = DataProcessor.get_value(record, selected_path) + except Exception: + value = None + + field_value = self._serialize_value(value) + + return { + "ui": { + # Important: send the plain list, not [list] + "field_choices": available_paths_list, + "selected_path": selected_path, + "available_paths": available_paths_text, + }, + "result": ( + field_value, + available_paths_text, + selected_path, + ), + } + + except Exception as exc: + logger.exception("ZMongoFieldSelector failure") + error_text = f"Error: {exc}" + return { + "ui": { + "field_choices": ["text"], + "selected_path": "text", + "available_paths": "", + }, + "result": ( + error_text, + "", + "text", + ), + } + +class ZMongoOperationsNode: + """Insert node using ZMongo.run_sync and SafeResult throughout.""" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "collection": ("STRING", {"default": "legal_codex"}), + "data_json": ("STRING", {"multiline": True}), + "operation_type": (["Add Unique (No Update)", "Standard Insert"],), + } + } + + RETURN_TYPES = ("STRING", "INT") + RETURN_NAMES = ("result_json", "count") + FUNCTION = "execute_op" + CATEGORY = "ZMongo/Operations" + + def _insert_unique_no_update(self, zmongo: ZMongo, collection: str, records: List[Dict[str, Any]]) -> SafeResult: + inserted_ids: List[Any] = [] + skipped = 0 + errors: List[str] = [] + + for record in records: + if not isinstance(record, dict): + errors.append("Record is not an object") + continue + + if "_id" in record: + exists = zmongo.run_sync(zmongo.find_one, collection, {"_id": record.get("_id")}) + if isinstance(exists, SafeResult) and exists.success and exists.data is not None: + skipped += 1 + continue + if isinstance(exists, SafeResult) and not exists.success: + errors.append(str(exists.error)) + continue + + result = zmongo.run_sync(zmongo.insert_one_async, collection, record) + if isinstance(result, SafeResult) and result.success: + inserted_ids.append((result.data or {}).get("inserted_id")) + elif isinstance(result, SafeResult): + if "duplicate" in str(result.error).lower(): + skipped += 1 + else: + errors.append(str(result.error)) + else: + errors.append(f"Unexpected insert result type: {type(result).__name__}") + + payload = { + "inserted_count": len(inserted_ids), + "inserted_ids": inserted_ids, + "skipped_count": skipped, + "error_count": len(errors), + "errors": errors, + } + return SafeResult.ok(payload) if not errors else SafeResult.fail("Some records failed", data=payload) + + def execute_op(self, collection, data_json, operation_type): + parsed = _coerce_records(data_json) + if not parsed.success: + return (_safe_json(parsed), 0) + + try: + zmongo = _get_zmongo() + records = parsed.data or [] + + if operation_type == "Add Unique (No Update)": + result = self._insert_unique_no_update(zmongo, collection, records) + count = ((result.data or {}).get("inserted_count", 0) if isinstance(result.data, dict) else 0) + return (_safe_json(result), int(count)) + + result = zmongo.run_sync(zmongo.insert_many, collection, records) + if not isinstance(result, SafeResult): + result = SafeResult.fail(f"Unexpected insert result type: {type(result).__name__}") + + inserted_ids = (result.data or {}).get("inserted_ids", []) if isinstance(result.data, dict) else [] + return (_safe_json(result), len(inserted_ids)) + except DuplicateKeyError as exc: + failure = SafeResult.fail(str(exc)) + return (_safe_json(failure), 0) + except Exception as exc: + logger.exception("ZMongoOperationsNode failure") + failure = SafeResult.fail(str(exc)) + return (_safe_json(failure), 0) + + +class ZMongoConfigNode: + """Configures the shared ZMongo instance used by the other nodes.""" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "mongo_uri": ("STRING", {"default": "mongodb://127.0.0.1:27017"}), + "db_name": ("STRING", {"default": "test"}), + } + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("status",) + FUNCTION = "configure" + CATEGORY = "ZMongo/Config" + + def configure(self, mongo_uri, db_name): + try: + zmongo = _get_zmongo(uri=mongo_uri, db_name=db_name) + ping = zmongo.run_sync(zmongo.sync_timestamp) + if isinstance(ping, SafeResult) and ping.success: + return (f"ZMongo connected to {db_name}",) + error_text = ping.error if isinstance(ping, SafeResult) else f"Unexpected result: {type(ping).__name__}" + return (f"ZMongo config failed: {error_text}",) + except Exception as exc: + logger.exception("ZMongoConfigNode failure") + return (f"ZMongo config failed: {exc}",) + + +class ZMongoTextFetcher: + """Fetches the `text` field from a document using ZMongo.run_sync.""" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "collection_name": ("STRING", {"default": "test"}), + "document_id": ("STRING", {"default": ""}), + }, + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("text_value",) + FUNCTION = "fetch_text" + CATEGORY = "ZMongo" + + def fetch_text(self, collection_name, document_id): + if not document_id or not document_id.strip(): + return ("No ID Provided",) + + try: + query_id = ObjectId(document_id) if ObjectId.is_valid(document_id) else document_id + zmongo = _get_zmongo() + result = zmongo.run_sync(zmongo.find_one, collection_name, {"_id": query_id}) + if not isinstance(result, SafeResult): + result = SafeResult.fail(f"Unexpected fetch result type: {type(result).__name__}") + + if not result.success: + return (f"Error: {result.error}",) + if not result.data: + return ("Document not found",) + + text_value = result.data.get("text") + if text_value is None: + return ("Field 'text' not found",) + if isinstance(text_value, (dict, list)): + return (_safe_json(text_value),) + return (str(text_value),) + except Exception as exc: + logger.exception("ZMongoTextFetcher failure") + return (f"Error: {exc}",) + + +NODE_CLASS_MAPPINGS = { + "ZMongoConfigNode": ZMongoConfigNode, + "ZMongoTextFetcher": ZMongoTextFetcher, + "ZMongoOperationsNode": ZMongoOperationsNode, + "ZMongoRecordSplitter": ZMongoRecordSplitter, + "ZMongoFieldSelector": ZMongoFieldSelector, + "ZRetrieverNode": ZRetrieverNode, + "ZMongoDatabaseBrowserNode": ZMongoDatabaseBrowserNode, + "ZMongoRecordLoopNode": ZMongoRecordLoopNode, + "ZMongoSaveTextNode": ZMongoSaveTextNode, + "ZMongoDataPassThroughNode": ZMongoDataPassThroughNode, + "ZMongoSaveBatchTextNode": ZMongoSaveBatchTextNode, + "ZMongoSaveValueNode": ZMongoSaveValueNode, + "ZMongoFlattenedFieldSelector": ZMongoFlattenedFieldSelector, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "ZMongoConfigNode": "ZMongo Config", + "ZMongoTextFetcher": "ZMongo Text Fetcher", + "ZMongoOperationsNode": "ZMongo Operations", + "ZMongoRecordSplitter": "ZMongo Record Splitter", + "ZMongoFieldSelector": "ZMongo Field Selector", + "ZRetrieverNode": "ZMongo Retriever", + "ZMongoDatabaseBrowserNode": "ZMongo Database Browser", + "ZMongoRecordLoopNode": "ZMongo Record Loop", + "ZMongoSaveTextNode": "ZMongo Save Text Node", + "ZMongoDataPassThroughNode": "ZMongo Data Pass Through Node", + "ZMongoSaveBatchTextNode": "ZMongo Save Batch Text Node", + "ZMongoSaveValueNode": "ZMongo Save Value Node", + "ZMongoFlattenedFieldSelector": "ZMongo Flattened Field Selector", +} \ No newline at end of file diff --git a/nodes/zmongo_record_editor_node.py b/nodes/zmongo_record_editor_node.py new file mode 100644 index 0000000..68390a5 --- /dev/null +++ b/nodes/zmongo_record_editor_node.py @@ -0,0 +1,259 @@ +import json +import logging +from typing import Any, Dict, List, Tuple + +from bson import ObjectId + +from .zmongo_toolbag.zmongo import ZMongo +from .zmongo_toolbag.data_processor import DataProcessor + +logger = logging.getLogger(__name__) + + +class ZMongoRecordEditorNode: + CATEGORY = "ZMongo" + FUNCTION = "get_record" + RETURN_TYPES = ("STRING", "STRING", "STRING", "STRING") + RETURN_NAMES = ( + "record_json", + "record_id", + "save_status", + "flattened_field_names_json", + ) + + @classmethod + def _get_zmongo(cls) -> ZMongo: + return ZMongo() + + @staticmethod + def _clean_record_id(value: Any) -> str: + raw = "" if value is None else str(value) + cleaned = raw.strip() + + if len(cleaned) >= 2 and cleaned[0] == cleaned[-1] and cleaned[0] in {"'", '"'}: + cleaned = cleaned[1:-1].strip() + + return cleaned + + @classmethod + def _build_id_candidates(cls, record_id: Any) -> List[Any]: + cleaned = cls._clean_record_id(record_id) + if not cleaned: + return [] + + candidates: List[Any] = [cleaned] + if ObjectId.is_valid(cleaned): + candidates.insert(0, ObjectId(cleaned)) + return candidates + + @staticmethod + def _stringify_value(value: Any) -> str: + if value is None: + return "" + if isinstance(value, str): + return value + try: + return json.dumps( + DataProcessor.to_json_compatible(value), + ensure_ascii=False, + indent=2, + default=str, + ) + except Exception: + return str(value) + + @classmethod + def get_collection_names(cls) -> List[str]: + zmongo = None + try: + zmongo = cls._get_zmongo() + result = zmongo.run_sync(zmongo.list_collections_async) + + if not result or not getattr(result, "success", False): + logger.warning( + "Could not load collection names: %s", + getattr(result, "error", None), + ) + return [""] + + data = result.data or {} + if not isinstance(data, dict): + return [""] + + names = data.get("collections", []) + if not isinstance(names, list): + return [""] + + names = [str(name) for name in names if str(name).strip()] + return names or [""] + + except Exception as exc: + logger.exception("Error loading collection names: %s", exc) + return [""] + finally: + if zmongo: + try: + zmongo.close() + except Exception: + pass + + @classmethod + def fetch_record( + cls, + collection_name: str, + record_id: str, + fallback_record_json: str = "", + ) -> Dict[str, Any]: + zmongo = None + try: + clean_id = cls._clean_record_id(record_id) + + if clean_id: + if not collection_name or collection_name.startswith("<"): + return {} + + zmongo = cls._get_zmongo() + + for candidate_id in cls._build_id_candidates(clean_id): + result = zmongo.find_one(collection_name, {"_id": candidate_id}) + + if result and getattr(result, "success", False): + doc = result.original() if hasattr(result, "original") else result.data + if isinstance(doc, dict) and doc: + return doc + + logger.warning( + "Could not fetch record '%s' from '%s' using string/ObjectId lookup.", + clean_id, + collection_name, + ) + return {} + + if fallback_record_json and str(fallback_record_json).strip(): + parsed = json.loads(fallback_record_json) + return parsed if isinstance(parsed, dict) else {} + + return {} + + except Exception as exc: + logger.exception( + "Error fetching record for collection='%s', record_id='%s': %s", + collection_name, + record_id, + exc, + ) + return {} + finally: + if zmongo: + try: + zmongo.close() + except Exception: + pass + + @classmethod + def get_flattened_record_pairs( + cls, + collection_name: str, + record_id: str, + fallback_record_json: str = "", + ) -> List[Tuple[str, str]]: + record = cls.fetch_record(collection_name, record_id, fallback_record_json) + if not record: + return [] + + flat = DataProcessor.flatten_json(record) + if not isinstance(flat, dict): + return [] + + return [ + (str(path), cls._stringify_value(value)) + for path, value in sorted(flat.items(), key=lambda item: str(item[0])) + ] + + @classmethod + def INPUT_TYPES(cls) -> Dict[str, Any]: + collections = cls.get_collection_names() + default_collection = collections[0] if collections else "" + + return { + "required": { + "collection_name": (collections, {"default": default_collection}), + "record_id": ("STRING", {"default": "", "multiline": False}), + }, + "optional": { + "selected_record_json": ("STRING", {"default": "", "forceInput": True}), + "save_status_in": ("STRING", {"default": "", "forceInput": True}), + }, + } + + @classmethod + def VALIDATE_INPUTS( + cls, + collection_name: str, + record_id: str, + selected_record_json: str = "", + save_status_in: str = "", + ): + return True + + @classmethod + def IS_CHANGED( + cls, + collection_name: str, + record_id: str, + selected_record_json: str = "", + save_status_in: str = "", + ): + return f"{collection_name}|{record_id}|{selected_record_json}|{save_status_in}" + + def get_record( + self, + collection_name: str, + record_id: str, + selected_record_json: str = "", + save_status_in: str = "", + ): + record = self.fetch_record( + collection_name=collection_name, + record_id=record_id, + fallback_record_json=selected_record_json, + ) + + if not record: + return ( + "{}", + self._clean_record_id(record_id), + str(save_status_in or ""), + json.dumps([], indent=2), + ) + + resolved_record_id = str(record.get("_id", self._clean_record_id(record_id))) + flattened_paths = [ + path + for path, _ in self.get_flattened_record_pairs( + collection_name=collection_name, + record_id=resolved_record_id, + fallback_record_json=selected_record_json, + ) + ] + + return ( + json.dumps( + DataProcessor.to_json_compatible(record), + ensure_ascii=False, + indent=2, + default=str, + ), + resolved_record_id, + str(save_status_in or ""), + json.dumps(flattened_paths, ensure_ascii=False, indent=2), + ) + + +NODE_CLASS_MAPPINGS = { + "ZMongoRecordEditorNode": ZMongoRecordEditorNode, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "ZMongoRecordEditorNode": "📝 ZMongo Record Editor", +} \ No newline at end of file diff --git a/nodes/zmongo_record_splitter.py b/nodes/zmongo_record_splitter.py new file mode 100644 index 0000000..9dc6a5d --- /dev/null +++ b/nodes/zmongo_record_splitter.py @@ -0,0 +1,55 @@ +import json +from .zmongo_toolbag.data_processor import DataProcessor + + +class ZMongoRecordSplitter: + """Converts JSON arrays into individual processable strings for batching.""" + + @classmethod + def INPUT_TYPES(cls): + return {"required": {"records_json": ("STRING", {"multiline": True})}} + + RETURN_TYPES = ("STRING", "INT") + RETURN_NAMES = ("record_json_list", "count") + OUTPUT_IS_LIST = (True, False) + FUNCTION = "split" + CATEGORY = "ZMongo/Database" + + def split(self, records_json): + try: + # Handles concatenated objects and arrays + data = json.loads(records_json) + records = data if isinstance(data, list) else [data] + output = [DataProcessor.to_json(r) for r in records] + return (output, len(output)) + except Exception as e: + return ([json.dumps({"error": str(e)})], 0) + + +class ZMongoFieldSelector: + """Dynamic dot-notation field extractor with metadata for UI dropdowns.""" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "record_json": ("STRING", {"forceInput": True}), + "field_path": ("STRING", {"default": "metadata.case_name"}), + } + } + + RETURN_TYPES = ("STRING", "STRING") + RETURN_NAMES = ("field_value", "available_paths") + FUNCTION = "select_field" + CATEGORY = "ZMongo/Database" + + def select_field(self, record_json, field_path): + record = json.loads(record_json) + # Extract nested values using DataProcessor dot-notation + val = DataProcessor.get_value(record, field_path) + + # Output paths to allow JS extension to build dropdowns + flattened = DataProcessor.flatten_dict(record) + available_paths = "\n".join(sorted(flattened.keys())) + + return (str(val), available_paths) \ No newline at end of file diff --git a/nodes/zmongo_tabular_record_view_node.py b/nodes/zmongo_tabular_record_view_node.py new file mode 100644 index 0000000..6b55b6d --- /dev/null +++ b/nodes/zmongo_tabular_record_view_node.py @@ -0,0 +1,367 @@ +import json +import logging +import re +from typing import Any, Dict, List, Tuple + +from .zmongo_toolbag.zmongo import ZMongo +from .zmongo_toolbag.data_processor import DataProcessor + +logger = logging.getLogger(__name__) + + +class ZMongoTabularRecordViewNode: + CATEGORY = "ZMongo" + FUNCTION = "select_record" + RETURN_TYPES = ("STRING", "STRING", "STRING", "INT", "STRING", "STRING") + RETURN_NAMES = ( + "selected_record_json", + "selected_record_id", + "selected_field_name", + "selected_record_index", + "selected_field_value", + "flattened_headings_json", + ) + + @staticmethod + @staticmethod + def _value_to_searchable_string(value: Any) -> str: + if value is None: + return "" + if isinstance(value, str): + return value + if isinstance(value, (bool, int, float)): + return str(value) + try: + return json.dumps( + DataProcessor.to_json_compatible(value), + ensure_ascii=False, + sort_keys=True, + default=str, + ) + except Exception: + return str(value) + + @staticmethod + def _wildcard_to_regex(pattern: str) -> str: + """ + Convert a simple wildcard pattern to regex. + Supported: + * => any characters + ? => any single character + """ + escaped = re.escape(pattern) + escaped = escaped.replace(r"\*", ".*") + escaped = escaped.replace(r"\?", ".") + return f"^{escaped}$" + + @classmethod + def _matches_search_text(cls, value: Any, search_text: str) -> bool: + """ + Matching rules: + - empty search_text => match everything + - if search_text contains * or ?, use wildcard matching + - otherwise do case-insensitive substring matching + """ + haystack = cls._value_to_searchable_string(value) + needle = str(search_text or "") + + if not needle: + return True + + haystack_lower = haystack.lower() + needle_lower = needle.lower() + + if "*" in needle or "?" in needle: + regex = cls._wildcard_to_regex(needle_lower) + return re.search(regex, haystack_lower) is not None + + return needle_lower in haystack_lower + + @classmethod + def get_filtered_table_payload( + cls, + collection_name: str, + search_text: str = "", + flattened_field_name: str = "", + ) -> Tuple[List[str], List[Dict[str, Any]], List[str]]: + records = cls.get_all_records(collection_name) + search_text = str(search_text or "").strip() + flattened_field_name = str(flattened_field_name or "").strip() + + headings: List[str] = [] + seen = set() + flat_records: List[Dict[str, Any]] = [] + record_ids: List[str] = [] + + for record in records: + flat = DataProcessor.flatten_json(record) + normalized_flat = { + str(k): DataProcessor.to_json_compatible(v) + for k, v in flat.items() + } + + if search_text: + if flattened_field_name: + target_value = normalized_flat.get(flattened_field_name) + matched = cls._matches_search_text(target_value, search_text) + else: + matched = any( + cls._matches_search_text(v, search_text) + for v in normalized_flat.values() + ) + + if not matched: + continue + + flat_records.append(normalized_flat) + + record_id = str(record.get("_id", "")) + record_ids.append(record_id) + + for key in normalized_flat.keys(): + if key not in seen: + seen.add(key) + headings.append(key) + + return headings, flat_records, record_ids + + @classmethod + def _get_zmongo(cls) -> ZMongo: + return ZMongo() + + @staticmethod + def _clamp_index(value: int, max_len: int) -> int: + if max_len <= 0: + return 0 + return max(0, min(int(value), max_len - 1)) + + @staticmethod + def _stringify_value(value: Any) -> str: + if value is None: + return "" + if isinstance(value, str): + return value + try: + return json.dumps( + DataProcessor.to_json_compatible(value), + ensure_ascii=False, + indent=2, + default=str, + ) + except Exception: + return str(value) + + @classmethod + def get_collection_names(cls) -> List[str]: + zmongo = None + try: + zmongo = cls._get_zmongo() + result = zmongo.run_sync(zmongo.list_collections_async) + + if not result or not getattr(result, "success", False): + logger.warning( + "Could not load collection names: %s", + getattr(result, "error", None), + ) + return [""] + + data = result.data or {} + if not isinstance(data, dict): + return [""] + + names = data.get("collections", []) + if not isinstance(names, list): + return [""] + + names = [str(name) for name in names if str(name).strip()] + return names or [""] + + except Exception as exc: + logger.exception("Error loading collection names: %s", exc) + return [""] + finally: + if zmongo: + try: + zmongo.close() + except Exception: + pass + + @classmethod + def get_all_records(cls, collection_name: str) -> List[Dict[str, Any]]: + zmongo = None + try: + if not collection_name or collection_name.startswith("<"): + return [] + + zmongo = cls._get_zmongo() + + result = zmongo.find_many( + collection_name, + query={}, + sort=[("_id", 1)], + limit=None, # all records + ) + + if not result or not getattr(result, "success", False): + logger.warning( + "Could not load records for '%s': %s", + collection_name, + getattr(result, "error", None), + ) + return [] + + data = result.original() if hasattr(result, "original") else result.data + if not isinstance(data, list): + return [] + + return [record for record in data if isinstance(record, dict)] + + except Exception as exc: + logger.exception("Error loading records for '%s': %s", collection_name, exc) + return [] + finally: + if zmongo: + try: + zmongo.close() + except Exception: + pass + + @classmethod + def get_table_payload(cls, collection_name: str) -> Tuple[List[str], List[Dict[str, Any]], List[str]]: + records = cls.get_all_records(collection_name) + + headings: List[str] = [] + seen = set() + flat_records: List[Dict[str, Any]] = [] + record_ids: List[str] = [] + + for record in records: + flat = DataProcessor.flatten_json(record) + normalized_flat = { + str(k): DataProcessor.to_json_compatible(v) + for k, v in flat.items() + } + flat_records.append(normalized_flat) + + record_id = str(record.get("_id", "")) + record_ids.append(record_id) + + for key in normalized_flat.keys(): + if key not in seen: + seen.add(key) + headings.append(key) + + return headings, flat_records, record_ids + + @classmethod + def INPUT_TYPES(cls) -> Dict[str, Any]: + collections = cls.get_collection_names() + default_collection = collections[0] if collections else "" + + return { + "required": { + "collection_name": (collections, {"default": default_collection}), + "flattened_field_name": ( + "STRING", + { + "default": "", + "multiline": False, + }, + ), + "record_id": ( + "STRING", + { + "default": "", + "multiline": False, + }, + ), + "selected_record_index": ( + "INT", + { + "default": 0, + "min": 0, + "max": 1000000, + "step": 1, + "display": "number", + }, + ), + } + } + + @classmethod + def VALIDATE_INPUTS(cls, collection_name, flattened_field_name, record_id, selected_record_index): + return True + + @classmethod + def IS_CHANGED(cls, collection_name: str, flattened_field_name: str, record_id: str, selected_record_index: int): + return f"{collection_name}|{flattened_field_name}|{record_id}|{selected_record_index}" + + def select_record( + self, + collection_name: str, + flattened_field_name: str, + record_id: str, + selected_record_index: int, + ): + records = self.get_all_records(collection_name) + headings, _, record_ids = self.get_table_payload(collection_name) + + if not records: + return ( + "{}", + "", + str(flattened_field_name or ""), + 0, + "", + json.dumps([], indent=2), + ) + + selected_idx = None + + if record_id: + record_id = str(record_id).strip() + for idx, rid in enumerate(record_ids): + if rid == record_id: + selected_idx = idx + break + + if selected_idx is None: + selected_idx = self._clamp_index(selected_record_index, len(records)) + + selected_record = records[selected_idx] + selected_record_id = str(selected_record.get("_id", "")) + + # default record_id behavior: if blank, use record 0 / selected record + if not record_id: + record_id = selected_record_id + + selected_field_name = str(flattened_field_name or "").strip() + selected_field_value = "" + if selected_field_name: + value = DataProcessor.get_value(selected_record, selected_field_name) + selected_field_value = self._stringify_value(value) + + selected_record_json = json.dumps( + DataProcessor.to_json_compatible(selected_record), + ensure_ascii=False, + indent=2, + default=str, + ) + + return ( + selected_record_json, + selected_record_id, + selected_field_name, + int(selected_idx), + selected_field_value, + json.dumps(headings, ensure_ascii=False, indent=2), + ) + + +NODE_CLASS_MAPPINGS = { + "ZMongoTabularRecordViewNode": ZMongoTabularRecordViewNode, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "ZMongoTabularRecordViewNode": "📊 ZMongo Tabular Record View", +} \ No newline at end of file diff --git a/web/js/app.js b/web/js/app.js new file mode 100644 index 0000000..6652405 --- /dev/null +++ b/web/js/app.js @@ -0,0 +1,169 @@ +import { app } from "/scripts/app.web"; + +function findWidget(node, name) { + return node.widgets?.find((w) => w.name === name); +} + +async function fetchFlattenedFields(collectionName) { + if (!collectionName || String(collectionName).startsWith("<")) { + return []; + } + + const response = await fetch( + `/zmongo/flattened_fields?collection_name=${encodeURIComponent(collectionName)}` + ); + + if (!response.ok) { + throw new Error(`HTTP ${response.status}`); + } + + const data = await response.json(); + if (!data.success || !Array.isArray(data.fields)) { + return []; + } + + return data.fields; +} + +function setComboValues(widget, values, preferredValue = null) { + const safeValues = values.length ? values : [""]; + + widget.options = widget.options || {}; + widget.options.values = safeValues; + + if (preferredValue && safeValues.includes(preferredValue)) { + widget.value = preferredValue; + } else if (!safeValues.includes(widget.value)) { + widget.value = safeValues[0]; + } +} + +function clampIndex(idx, len) { + if (!len || len <= 0) return 0; + if (Number.isNaN(idx)) return 0; + return Math.max(0, Math.min(idx, len - 1)); +} + +app.registerExtension({ + name: "ZMongo.FlattenedFieldDropdown", + + async nodeCreated(node) { + if (node.comfyClass !== "ZMongoFlattenedFieldDropdownNode") { + return; + } + + const collectionWidget = findWidget(node, "collection_name"); + const fieldWidget = findWidget(node, "field_name"); + const indexWidget = findWidget(node, "field_index"); + + if (!collectionWidget || !fieldWidget || !indexWidget) { + return; + } + + let isUpdating = false; + + const syncIndexFromField = () => { + const values = fieldWidget.options?.values || []; + const idx = values.indexOf(fieldWidget.value); + indexWidget.value = idx >= 0 ? idx : 0; + node.setDirtyCanvas(true, true); + }; + + const syncFieldFromIndex = () => { + const values = fieldWidget.options?.values || []; + if (!values.length) { + fieldWidget.value = ""; + indexWidget.value = 0; + node.setDirtyCanvas(true, true); + return; + } + + let idx = Number(indexWidget.value ?? 0); + idx = clampIndex(idx, values.length); + indexWidget.value = idx; + fieldWidget.value = values[idx]; + node.setDirtyCanvas(true, true); + }; + + const refreshFields = async ({ preserveField = true } = {}) => { + if (isUpdating) return; + isUpdating = true; + + try { + const collectionName = collectionWidget.value; + const oldFieldValue = fieldWidget.value; + const fields = await fetchFlattenedFields(collectionName); + + let preferredValue = null; + if ( + preserveField && + oldFieldValue && + !String(oldFieldValue).startsWith("<") && + fields.includes(oldFieldValue) + ) { + preferredValue = oldFieldValue; + } + + setComboValues(fieldWidget, fields, preferredValue); + + const values = fieldWidget.options?.values || []; + let idx = Number(indexWidget.value ?? 0); + idx = clampIndex(idx, values.length); + indexWidget.value = idx; + + if (!preferredValue) { + fieldWidget.value = values[idx] ?? ""; + } + + syncIndexFromField(); + node.setDirtyCanvas(true, true); + } catch (err) { + console.error("ZMongo flattened field dropdown refresh failed:", err); + setComboValues(fieldWidget, [""]); + fieldWidget.value = ""; + indexWidget.value = 0; + node.setDirtyCanvas(true, true); + } finally { + isUpdating = false; + } + }; + + const originalCollectionCallback = collectionWidget.callback; + collectionWidget.callback = async (...args) => { + if (originalCollectionCallback) { + originalCollectionCallback.apply(collectionWidget, args); + } + await refreshFields({ preserveField: false }); + }; + + const originalFieldCallback = fieldWidget.callback; + fieldWidget.callback = (...args) => { + if (originalFieldCallback) { + originalFieldCallback.apply(fieldWidget, args); + } + syncIndexFromField(); + }; + + const originalIndexCallback = indexWidget.callback; + indexWidget.callback = (...args) => { + if (originalIndexCallback) { + originalIndexCallback.apply(indexWidget, args); + } + syncFieldFromIndex(); + }; + + // Initial population from the default selected collection. + await refreshFields({ preserveField: false }); + + // Run again on next tick because ComfyUI sometimes finishes widget restoration + // after nodeCreated, which can otherwise leave the placeholder stuck. + setTimeout(async () => { + await refreshFields({ preserveField: false }); + }, 0); + + // One more delayed refresh for restored workflows / startup timing issues. + setTimeout(async () => { + await refreshFields({ preserveField: false }); + }, 150); + }, +}); \ No newline at end of file diff --git a/web/js/preset_panel.js b/web/js/preset_panel.js new file mode 100644 index 0000000..b8c22c9 --- /dev/null +++ b/web/js/preset_panel.js @@ -0,0 +1,336 @@ +import { app } from "../../scripts/app.web"; + +const EXT_ID = "ZMongo.PresetPanel"; +const PANEL_ID = "zmongo-preset-panel"; + +function getWorkflowNodes() { + return app.graph?._nodes ?? []; +} + +function getNodeTitle(node) { + return String(node?.title || "").trim(); +} + +function findNodeByTitle(title) { + return getWorkflowNodes().find((node) => getNodeTitle(node) === title) || null; +} + +function findWidget(node, widgetName) { + return node?.widgets?.find((w) => w?.name === widgetName) || null; +} + +function coerceValueForWidget(widget, value) { + if (!widget) return value; + const current = widget.value; + if (typeof current === "number" && typeof value === "string") { + const n = Number(value); + return Number.isNaN(n) ? value : n; + } + return value; +} + +function setWidgetValue(node, widgetName, value) { + const widget = findWidget(node, widgetName); + if (!widget) return false; + + const finalValue = coerceValueForWidget(widget, value); + widget.value = finalValue; + + if (typeof widget.callback === "function") { + try { + widget.callback(finalValue); + } catch {} + } + + if (typeof node.onWidgetChanged === "function") { + try { + node.onWidgetChanged(widgetName, finalValue, widget); + } catch {} + } + + node.setDirtyCanvas?.(true, true); + return true; +} + +function applyPreset(preset) { + const failures = []; + const updates = preset?.updates ?? {}; + + for (const [nodeTitle, widgetValues] of Object.entries(updates)) { + const node = findNodeByTitle(nodeTitle); + if (!node) { + failures.push(`Missing node: ${nodeTitle}`); + continue; + } + + for (const [widgetName, value] of Object.entries(widgetValues)) { + const ok = setWidgetValue(node, widgetName, value); + if (!ok) failures.push(`Missing widget: ${nodeTitle}.${widgetName}`); + } + } + + app.graph?.setDirtyCanvas?.(true, true); + + if (failures.length) { + alert(`Preset applied with warnings:\n\n${failures.join("\n")}`); + } +} + +function queuePrompt() { + try { + app.queuePrompt?.(0); + } catch (err) { + console.warn("[ZMongo PresetPanel] Queue failed", err); + } +} + +async function loadPresets(workflow = "default") { + const res = await fetch(`/zmongo/presets?workflow=${encodeURIComponent(workflow)}`); + if (!res.ok) { + throw new Error(`Failed to load presets (${res.status})`); + } + const data = await res.json(); + if (!data.success) { + throw new Error(data.error || "Failed to load presets"); + } + + const groups = {}; + for (const p of data.presets ?? []) { + const g = p.group || "Default"; + if (!groups[g]) groups[g] = []; + groups[g].push(p); + } + + return { + groups: Object.entries(groups).map(([name, presets]) => ({ name, presets })) + }; +} + +function extractCurrentValues(nodeTitles) { + const out = {}; + for (const title of nodeTitles) { + const node = findNodeByTitle(title); + if (!node || !node.widgets) continue; + + out[title] = {}; + for (const widget of node.widgets) { + if (!widget?.name) continue; + out[title][widget.name] = widget.value; + } + } + return out; +} + +async function savePreset(workflow, name, group, updates) { + const res = await fetch("/zmongo/presets", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ + workflow, + name, + group, + updates + }) + }); + + const data = await res.json(); + if (!res.ok || !data.success) { + throw new Error(data.error || "Failed to save preset"); + } +} + +async function deletePreset(workflow, name) { + const res = await fetch(`/zmongo/presets/${encodeURIComponent(workflow)}/${encodeURIComponent(name)}`, { + method: "DELETE" + }); + + const data = await res.json(); + if (!res.ok || !data.success) { + throw new Error(data.error || "Failed to delete preset"); + } +} + +function makeButton(label, onClick) { + const btn = document.createElement("button"); + btn.textContent = label; + btn.style.width = "100%"; + btn.style.marginBottom = "8px"; + btn.style.padding = "8px"; + btn.style.borderRadius = "6px"; + btn.style.cursor = "pointer"; + btn.onclick = onClick; + return btn; +} + +function makeSectionTitle(text) { + const el = document.createElement("div"); + el.textContent = text; + el.style.fontWeight = "bold"; + el.style.margin = "12px 0 8px"; + el.style.fontSize = "14px"; + return el; +} + +function makeInfo(text) { + const el = document.createElement("div"); + el.textContent = text; + el.style.fontSize = "12px"; + el.style.opacity = "0.8"; + el.style.marginBottom = "10px"; + return el; +} + +async function buildPanel(container, workflow = "default") { + container.innerHTML = ""; + container.appendChild(makeInfo(`Workflow key: ${workflow}`)); + + const config = await loadPresets(workflow); + + for (const group of config.groups ?? []) { + container.appendChild(makeSectionTitle(group.name)); + + for (const preset of group.presets ?? []) { + const row = document.createElement("div"); + row.style.display = "flex"; + row.style.gap = "6px"; + row.style.marginBottom = "6px"; + + const applyBtn = makeButton(preset.name, () => applyPreset(preset)); + applyBtn.style.flex = "1"; + + const applyQueueBtn = makeButton("Apply + Queue", () => { + applyPreset(preset); + queuePrompt(); + }); + applyQueueBtn.style.flex = "1"; + + const delBtn = makeButton("Del", async () => { + if (!confirm(`Delete preset "${preset.name}"?`)) return; + try { + await deletePreset(workflow, preset.name); + await buildPanel(container, workflow); + } catch (err) { + alert(String(err)); + } + }); + delBtn.style.flex = "0 0 64px"; + + row.appendChild(applyBtn); + row.appendChild(applyQueueBtn); + row.appendChild(delBtn); + container.appendChild(row); + } + } + + container.appendChild(makeSectionTitle("Save Current Preset")); + + const nameInput = document.createElement("input"); + nameInput.placeholder = "Preset name"; + nameInput.style.width = "100%"; + nameInput.style.marginBottom = "6px"; + nameInput.style.padding = "8px"; + + const groupInput = document.createElement("input"); + groupInput.placeholder = "Group name"; + groupInput.value = "Default"; + groupInput.style.width = "100%"; + groupInput.style.marginBottom = "6px"; + groupInput.style.padding = "8px"; + + const nodesInput = document.createElement("textarea"); + nodesInput.placeholder = "Node titles, one per line"; + nodesInput.style.width = "100%"; + nodesInput.style.minHeight = "120px"; + nodesInput.style.marginBottom = "6px"; + nodesInput.style.padding = "8px"; + nodesInput.value = [ + "MAIN_MODEL", + "MAIN_LORA", + "MAIN_SAMPLER", + "IMAGE_SETTINGS", + "VIDEO_MODEL", + "VIDEO_UNET", + "VIDEO_LORA", + "VIDEO_SAMPLER", + "VIDEO_SETTINGS" + ].join("\n"); + + container.appendChild(nameInput); + container.appendChild(groupInput); + container.appendChild(nodesInput); + + container.appendChild( + makeButton("Save Current Preset to ZMongo", async () => { + const name = nameInput.value.trim(); + const group = groupInput.value.trim() || "Default"; + const nodeTitles = nodesInput.value + .split("\n") + .map((s) => s.trim()) + .filter(Boolean); + + if (!name) { + alert("Enter a preset name."); + return; + } + + const updates = extractCurrentValues(nodeTitles); + + try { + await savePreset(workflow, name, group, updates); + await buildPanel(container, workflow); + } catch (err) { + alert(String(err)); + } + }) + ); + + container.appendChild(makeSectionTitle("Utility")); + + container.appendChild(makeButton("Queue Prompt", () => queuePrompt())); + + container.appendChild( + makeButton("Refresh Presets", async () => { + try { + await buildPanel(container, workflow); + } catch (err) { + alert(String(err)); + } + }) + ); +} + +app.registerExtension({ + name: EXT_ID, + + async setup() { + const panel = document.createElement("div"); + panel.id = PANEL_ID; + panel.style.padding = "10px"; + panel.style.height = "100%"; + panel.style.overflowY = "auto"; + panel.style.boxSizing = "border-box"; + + const workflow = "default"; + + try { + await buildPanel(panel, workflow); + } catch (err) { + panel.textContent = String(err); + console.error(err); + } + + if (app.extensionManager?.registerSidebarTab) { + app.extensionManager.registerSidebarTab({ + id: PANEL_ID, + icon: "list", + title: "Presets", + tooltip: "ZMongo preset panel", + type: "custom", + render: (el) => { + el.innerHTML = ""; + el.appendChild(panel); + } + }); + } + } +}); \ No newline at end of file diff --git a/web/js/zmongo_field_selector.js b/web/js/zmongo_field_selector.js new file mode 100644 index 0000000..67168f5 --- /dev/null +++ b/web/js/zmongo_field_selector.js @@ -0,0 +1,148 @@ +import { app } from "/scripts/app.web"; + +function findWidget(node, name) { + return node.widgets?.find((w) => w.name === name); +} + +async function fetchFlattenedFields(collectionName) { + if (!collectionName || String(collectionName).startsWith("<")) { + return []; + } + + const url = `/zmongo/flattened_fields?collection_name=${encodeURIComponent(collectionName)}`; + const response = await fetch(url); + + if (!response.ok) { + throw new Error(`HTTP ${response.status}`); + } + + const data = await response.json(); + if (!data.success || !Array.isArray(data.fields)) { + return []; + } + + return data.fields; +} + +function setWidgetOptions(widget, values) { + const safeValues = values.length ? values : [""]; + + widget.options = widget.options || {}; + widget.options.values = safeValues; + + if (!safeValues.includes(widget.value)) { + widget.value = safeValues[0]; + } +} + +function clampIndex(value, maxLen) { + if (!maxLen || maxLen <= 0) { + return 0; + } + const n = Number(value); + if (Number.isNaN(n)) { + return 0; + } + return Math.max(0, Math.min(n, maxLen - 1)); +} + +async function refreshFieldWidget(node, collectionWidget, fieldWidget, indexWidget) { + const collectionName = collectionWidget.value; + const fields = await fetchFlattenedFields(collectionName); + + setWidgetOptions(fieldWidget, fields); + + const values = fieldWidget.options?.values || []; + const idx = clampIndex(indexWidget.value, values.length); + + indexWidget.value = idx; + fieldWidget.value = values[idx] ?? ""; + + node.setDirtyCanvas(true, true); +} + +app.registerExtension({ + name: "ZMongo.FlattenedFieldDropdown", + + async nodeCreated(node) { + if (node.comfyClass !== "ZMongoFlattenedFieldDropdownNode") { + return; + } + + const collectionWidget = findWidget(node, "collection_name"); + const fieldWidget = findWidget(node, "field_name"); + const indexWidget = findWidget(node, "field_index"); + + if (!collectionWidget || !fieldWidget || !indexWidget) { + console.warn("ZMongo node widgets not found."); + return; + } + + const syncIndexFromField = () => { + const values = fieldWidget.options?.values || []; + const idx = values.indexOf(fieldWidget.value); + indexWidget.value = idx >= 0 ? idx : 0; + node.setDirtyCanvas(true, true); + }; + + const syncFieldFromIndex = () => { + const values = fieldWidget.options?.values || []; + if (!values.length) { + fieldWidget.value = ""; + indexWidget.value = 0; + node.setDirtyCanvas(true, true); + return; + } + + const idx = clampIndex(indexWidget.value, values.length); + indexWidget.value = idx; + fieldWidget.value = values[idx]; + node.setDirtyCanvas(true, true); + }; + + const originalCollectionCallback = collectionWidget.callback; + collectionWidget.callback = async (...args) => { + if (originalCollectionCallback) { + originalCollectionCallback.apply(collectionWidget, args); + } + + try { + await refreshFieldWidget(node, collectionWidget, fieldWidget, indexWidget); + } catch (err) { + console.error("Failed refreshing fields after collection change:", err); + } + }; + + const originalFieldCallback = fieldWidget.callback; + fieldWidget.callback = (...args) => { + if (originalFieldCallback) { + originalFieldCallback.apply(fieldWidget, args); + } + syncIndexFromField(); + }; + + const originalIndexCallback = indexWidget.callback; + indexWidget.callback = (...args) => { + if (originalIndexCallback) { + originalIndexCallback.apply(indexWidget, args); + } + syncFieldFromIndex(); + }; + + // initial startup refresh from default selected collection + try { + await refreshFieldWidget(node, collectionWidget, fieldWidget, indexWidget); + } catch (err) { + console.error("Initial ZMongo field refresh failed:", err); + } + + // delayed retry to handle restored workflow timing + setTimeout(async () => { + try { + await refreshFieldWidget(node, collectionWidget, fieldWidget, indexWidget); + } catch (err) { + console.error("Delayed ZMongo field refresh failed:", err); + } + }, 200); + }, +}); \ No newline at end of file diff --git a/web/js/zmongo_field_selector_dropdown.js b/web/js/zmongo_field_selector_dropdown.js new file mode 100644 index 0000000..c402ed9 --- /dev/null +++ b/web/js/zmongo_field_selector_dropdown.js @@ -0,0 +1,145 @@ +import { app } from "/scripts/app.web"; + +function findWidget(node, name) { + return node.widgets?.find((w) => w.name === name); +} + +async function fetchFlattenedFields(collectionName) { + if (!collectionName || String(collectionName).startsWith("<")) { + return []; + } + + const url = `/zmongo/flattened_fields?collection_name=${encodeURIComponent(collectionName)}`; + const response = await fetch(url); + + if (!response.ok) { + throw new Error(`HTTP ${response.status}`); + } + + const data = await response.json(); + if (!data.success || !Array.isArray(data.fields)) { + return []; + } + + return data.fields; +} + +function setWidgetOptions(widget, values) { + const safeValues = values.length ? values : [""]; + widget.options = widget.options || {}; + widget.options.values = safeValues; + + if (!safeValues.includes(widget.value)) { + widget.value = safeValues[0]; + } +} + +function clampIndex(value, maxLen) { + if (!maxLen || maxLen <= 0) { + return 0; + } + const n = Number(value); + if (Number.isNaN(n)) { + return 0; + } + return Math.max(0, Math.min(n, maxLen - 1)); +} + +async function refreshFieldWidget(node, collectionWidget, fieldWidget, indexWidget) { + const collectionName = collectionWidget.value; + const fields = await fetchFlattenedFields(collectionName); + + setWidgetOptions(fieldWidget, fields); + + const values = fieldWidget.options?.values || []; + const idx = clampIndex(indexWidget.value, values.length); + + indexWidget.value = idx; + fieldWidget.value = values[idx] ?? ""; + + node.setDirtyCanvas(true, true); +} + +app.registerExtension({ + name: "ZMongo.FlattenedFieldDropdown", + + async nodeCreated(node) { + if (node.comfyClass !== "ZMongoFlattenedFieldDropdownNode") { + return; + } + + const collectionWidget = findWidget(node, "collection_name"); + const fieldWidget = findWidget(node, "field_name"); + const indexWidget = findWidget(node, "field_index"); + + if (!collectionWidget || !fieldWidget || !indexWidget) { + console.warn("ZMongo node widgets not found."); + return; + } + + const syncIndexFromField = () => { + const values = fieldWidget.options?.values || []; + const idx = values.indexOf(fieldWidget.value); + indexWidget.value = idx >= 0 ? idx : 0; + node.setDirtyCanvas(true, true); + }; + + const syncFieldFromIndex = () => { + const values = fieldWidget.options?.values || []; + if (!values.length) { + fieldWidget.value = ""; + indexWidget.value = 0; + node.setDirtyCanvas(true, true); + return; + } + + const idx = clampIndex(indexWidget.value, values.length); + indexWidget.value = idx; + fieldWidget.value = values[idx]; + node.setDirtyCanvas(true, true); + }; + + const originalCollectionCallback = collectionWidget.callback; + collectionWidget.callback = async (...args) => { + if (originalCollectionCallback) { + originalCollectionCallback.apply(collectionWidget, args); + } + + try { + await refreshFieldWidget(node, collectionWidget, fieldWidget, indexWidget); + } catch (err) { + console.error("Failed refreshing fields after collection change:", err); + } + }; + + const originalFieldCallback = fieldWidget.callback; + fieldWidget.callback = (...args) => { + if (originalFieldCallback) { + originalFieldCallback.apply(fieldWidget, args); + } + syncIndexFromField(); + }; + + const originalIndexCallback = indexWidget.callback; + indexWidget.callback = (...args) => { + if (originalIndexCallback) { + originalIndexCallback.apply(indexWidget, args); + } + syncFieldFromIndex(); + }; + + try { + await refreshFieldWidget(node, collectionWidget, fieldWidget, indexWidget); + } catch (err) { + console.error("Initial ZMongo field refresh failed:", err); + } + + setTimeout(async () => { + try { + await refreshFieldWidget(node, collectionWidget, fieldWidget, indexWidget); + } catch (err) { + console.error("Delayed ZMongo field refresh failed:", err); + } + }, 200); + }, +}); \ No newline at end of file diff --git a/web/js/zmongo_record_editor.js b/web/js/zmongo_record_editor.js new file mode 100644 index 0000000..953c68d --- /dev/null +++ b/web/js/zmongo_record_editor.js @@ -0,0 +1,332 @@ +import { app } from "/scripts/app.web"; + +function findWidget(node, name) { + return node.widgets?.find((w) => w.name === name); +} + +async function fetchJson(url) { + const response = await fetch(url); + if (!response.ok) { + throw new Error(`HTTP ${response.status}`); + } + const data = await response.json(); + if (!data.success) { + throw new Error(data.error || "Unknown error"); + } + return data; +} + +async function postJson(url, body) { + const response = await fetch(url, { + method: "POST", + headers: { + "Content-Type": "application/json", + }, + body: JSON.stringify(body), + }); + + if (!response.ok) { + throw new Error(`HTTP ${response.status}`); + } + + const data = await response.json(); + if (!data.success) { + throw new Error(data.error || "Unknown error"); + } + + return data; +} + +async function loadRecordEditorPayload(collectionName, recordId, selectedRecordJson) { + const url = + `/zmongo/record_editor/load?collection_name=${encodeURIComponent(collectionName || "")}` + + `&record_id=${encodeURIComponent(recordId || "")}` + + `&selected_record_json=${encodeURIComponent(selectedRecordJson || "")}`; + + return await fetchJson(url); +} + +app.registerExtension({ + name: "ZMongo.RecordEditor", + + async nodeCreated(node) { + if (node.comfyClass !== "ZMongoRecordEditorNode") { + return; + } + + const collectionWidget = findWidget(node, "collection_name"); + const recordIdWidget = findWidget(node, "record_id"); + const selectedRecordJsonWidget = findWidget(node, "selected_record_json"); + const saveStatusWidget = findWidget(node, "save_status_in"); + + if (!collectionWidget || !recordIdWidget) { + console.warn("ZMongoRecordEditorNode widgets not found."); + return; + } + + const container = document.createElement("div"); + container.style.display = "flex"; + container.style.flexDirection = "column"; + container.style.gap = "8px"; + container.style.width = "100%"; + container.style.minHeight = "320px"; + container.style.maxHeight = "620px"; + container.style.background = "#111"; + container.style.border = "1px solid #333"; + container.style.borderRadius = "8px"; + container.style.padding = "8px"; + container.style.boxSizing = "border-box"; + + const infoBar = document.createElement("div"); + infoBar.style.fontSize = "12px"; + infoBar.style.color = "#aaa"; + infoBar.textContent = "Load a record to edit."; + + const toolbar = document.createElement("div"); + toolbar.style.display = "flex"; + toolbar.style.gap = "8px"; + toolbar.style.alignItems = "center"; + toolbar.style.flexWrap = "wrap"; + + const loadButton = document.createElement("button"); + loadButton.textContent = "Load Record"; + loadButton.style.padding = "6px 10px"; + loadButton.style.borderRadius = "6px"; + loadButton.style.border = "1px solid #555"; + loadButton.style.background = "#2a2a2a"; + loadButton.style.color = "#eee"; + loadButton.style.cursor = "pointer"; + + const saveButton = document.createElement("button"); + saveButton.textContent = "Save"; + saveButton.style.padding = "6px 10px"; + saveButton.style.borderRadius = "6px"; + saveButton.style.border = "1px solid #555"; + saveButton.style.background = "#2a2a2a"; + saveButton.style.color = "#eee"; + saveButton.style.cursor = "pointer"; + + const saveStatusLabel = document.createElement("div"); + saveStatusLabel.style.fontSize = "11px"; + saveStatusLabel.style.color = "#8fb3ff"; + saveStatusLabel.textContent = ""; + + toolbar.appendChild(loadButton); + toolbar.appendChild(saveButton); + + const editorWrap = document.createElement("div"); + editorWrap.style.flex = "1"; + editorWrap.style.overflow = "auto"; + editorWrap.style.border = "1px solid #333"; + editorWrap.style.borderRadius = "6px"; + editorWrap.style.background = "#181818"; + editorWrap.style.padding = "8px"; + editorWrap.style.display = "flex"; + editorWrap.style.flexDirection = "column"; + editorWrap.style.gap = "8px"; + + container.appendChild(infoBar); + container.appendChild(toolbar); + container.appendChild(saveStatusLabel); + container.appendChild(editorWrap); + + if (node.addDOMWidget) { + node.addDOMWidget("record_editor", "div", container, { + serialize: false, + hideOnZoom: false, + }); + } + + let currentRecord = {}; + let currentFields = []; + let fieldEditors = new Map(); + + function getSelectedRecordJsonValue() { + if (!selectedRecordJsonWidget) { + return ""; + } + return String(selectedRecordJsonWidget.value || ""); + } + + function getRecordIdValue() { + return String(recordIdWidget.value || "").trim(); + } + + function buildFieldEditor(path, value) { + const row = document.createElement("div"); + row.style.display = "flex"; + row.style.flexDirection = "column"; + row.style.gap = "4px"; + row.style.padding = "6px"; + row.style.border = "1px solid #2f2f2f"; + row.style.borderRadius = "6px"; + row.style.background = "#141414"; + + const label = document.createElement("div"); + label.textContent = path; + label.title = path; + label.style.fontSize = "12px"; + label.style.fontWeight = "600"; + label.style.color = "#cfd8ff"; + label.style.wordBreak = "break-word"; + + const editor = document.createElement("textarea"); + editor.value = value ?? ""; + editor.style.width = "100%"; + editor.style.minHeight = "48px"; + editor.style.maxHeight = "220px"; + editor.style.resize = "vertical"; + editor.style.padding = "6px 8px"; + editor.style.borderRadius = "6px"; + editor.style.border = "1px solid #444"; + editor.style.background = "#1b1b1b"; + editor.style.color = "#ddd"; + editor.style.whiteSpace = "pre-wrap"; + editor.style.overflowWrap = "anywhere"; + editor.style.boxSizing = "border-box"; + + row.appendChild(label); + row.appendChild(editor); + + return { row, editor }; + } + + function renderEditors(fields) { + editorWrap.innerHTML = ""; + fieldEditors = new Map(); + + if (!Array.isArray(fields) || !fields.length) { + const empty = document.createElement("div"); + empty.textContent = "No editable fields loaded."; + empty.style.color = "#888"; + empty.style.fontSize = "12px"; + editorWrap.appendChild(empty); + return; + } + + for (const field of fields) { + const path = String(field.path || ""); + const value = String(field.value ?? ""); + const { row, editor } = buildFieldEditor(path, value); + fieldEditors.set(path, editor); + editorWrap.appendChild(row); + } + } + + function collectChanges() { + const changes = {}; + for (const [path, editor] of fieldEditors.entries()) { + changes[path] = editor.value; + } + return changes; + } + + async function loadRecord() { + const collectionName = String(collectionWidget.value || "").trim(); + const recordId = getRecordIdValue(); + const selectedRecordJson = getSelectedRecordJsonValue(); + + const payload = await loadRecordEditorPayload( + collectionName, + recordId, + selectedRecordJson + ); + + currentRecord = payload.record || {}; + currentFields = Array.isArray(payload.fields) ? payload.fields : []; + + if (payload.record_id) { + recordIdWidget.value = payload.record_id; + } + + renderEditors(currentFields); + + const shownRecordId = payload.record_id || recordId || ""; + infoBar.textContent = shownRecordId + ? `Editing ${collectionName} / ${shownRecordId}` + : `Editing ${collectionName}`; + + saveStatusLabel.textContent = ""; + node.setDirtyCanvas(true, true); + } + + async function saveRecord() { + const collectionName = String(collectionWidget.value || "").trim(); + const recordId = getRecordIdValue(); + + if (!collectionName) { + throw new Error("collection_name is required"); + } + if (!recordId) { + throw new Error("record_id is required"); + } + + const changes = collectChanges(); + + const payload = await postJson("/zmongo/record_editor/save", { + collection_name: collectionName, + record_id: recordId, + changes, + }); + + currentRecord = payload.record || {}; + saveStatusLabel.textContent = `Saved ${collectionName} / ${recordId}`; + + if (saveStatusWidget) { + saveStatusWidget.value = `Saved ${collectionName} / ${recordId}`; + } + + await loadRecord(); + node.setDirtyCanvas(true, true); + } + + loadButton.addEventListener("click", async () => { + try { + await loadRecord(); + } catch (err) { + console.error("Failed to load record editor:", err); + infoBar.textContent = `Error: ${err.message}`; + } + }); + + saveButton.addEventListener("click", async () => { + try { + await saveRecord(); + } catch (err) { + console.error("Failed to save record editor:", err); + saveStatusLabel.textContent = `Error: ${err.message}`; + } + }); + + const originalCollectionCallback = collectionWidget.callback; + collectionWidget.callback = async (...args) => { + if (originalCollectionCallback) { + originalCollectionCallback.apply(collectionWidget, args); + } + try { + await loadRecord(); + } catch (err) { + console.error("Failed to reload editor after collection change:", err); + } + }; + + const originalRecordIdCallback = recordIdWidget.callback; + recordIdWidget.callback = async (...args) => { + if (originalRecordIdCallback) { + originalRecordIdCallback.apply(recordIdWidget, args); + } + try { + await loadRecord(); + } catch (err) { + console.error("Failed to reload editor after record_id change:", err); + } + }; + + try { + await loadRecord(); + } catch (err) { + console.error("Initial record editor load failed:", err); + infoBar.textContent = `Error: ${err.message}`; + } + }, +}); \ No newline at end of file diff --git a/web/js/zmongo_tabular_record_view.js b/web/js/zmongo_tabular_record_view.js new file mode 100644 index 0000000..f928ee2 --- /dev/null +++ b/web/js/zmongo_tabular_record_view.js @@ -0,0 +1,508 @@ +import { app } from "/scripts/app.web"; + +function findWidget(node, name) { + return node.widgets?.find((w) => w.name === name); +} + +async function fetchJson(url) { + const response = await fetch(url); + if (!response.ok) { + throw new Error(`HTTP ${response.status}`); + } + + const data = await response.json(); + if (!data.success) { + throw new Error(data.error || "Unknown error"); + } + + return data; +} + +async function fetchTabularRecords(collectionName) { + if (!collectionName || String(collectionName).startsWith("<")) { + return { + headings: [], + flat_records: [], + record_ids: [], + default_record_id: "", + record_count: 0, + }; + } + + const url = `/zmongo/tabular_records?collection_name=${encodeURIComponent(collectionName)}`; + return await fetchJson(url); +} + +async function fetchTabularRecordsSearch(collectionName, searchText, flattenedFieldName) { + if (!collectionName || String(collectionName).startsWith("<")) { + return { + headings: [], + flat_records: [], + record_ids: [], + default_record_id: "", + record_count: 0, + }; + } + + const url = + `/zmongo/tabular_records_search?collection_name=${encodeURIComponent(collectionName)}` + + `&search_text=${encodeURIComponent(searchText ?? "")}` + + `&flattened_field_name=${encodeURIComponent(flattenedFieldName ?? "")}`; + + return await fetchJson(url); +} + +function clampIndex(value, maxLen) { + if (!maxLen || maxLen <= 0) { + return 0; + } + + const n = Number(value); + if (Number.isNaN(n)) { + return 0; + } + + return Math.max(0, Math.min(n, maxLen - 1)); +} + +function buildCell(content, isHeader = false) { + const cell = document.createElement(isHeader ? "th" : "td"); + cell.style.border = "1px solid #444"; + cell.style.padding = "0"; + cell.style.verticalAlign = "top"; + cell.style.background = isHeader ? "#2a2a2a" : "#1d1d1d"; + cell.style.color = "#ddd"; + + const inner = document.createElement("div"); + inner.textContent = content ?? ""; + inner.style.minWidth = "140px"; + inner.style.maxWidth = "480px"; + inner.style.padding = "6px"; + inner.style.whiteSpace = "pre-wrap"; + inner.style.overflowWrap = "anywhere"; + inner.style.wordBreak = "break-word"; + inner.style.resize = "horizontal"; + inner.style.overflow = "auto"; + inner.style.boxSizing = "border-box"; + + if (isHeader) { + inner.style.fontWeight = "600"; + inner.style.cursor = "pointer"; + inner.style.userSelect = "none"; + } + + cell.appendChild(inner); + return { cell, inner }; +} + +app.registerExtension({ + name: "ZMongo.TabularRecordView", + + async nodeCreated(node) { + if (node.comfyClass !== "ZMongoTabularRecordViewNode") { + return; + } + + function updateSearchModeLabel() { + const flattenedFieldName = String(fieldWidget.value || "").trim(); + const searchText = String(exactSearchInput.value || "").trim(); + + const targetText = flattenedFieldName + ? `field "${flattenedFieldName}"` + : "all fields"; + + const queryText = searchText + ? ` | query: "${searchText}"` + : ""; + + searchModeLabel.textContent = `Search target: ${targetText}${queryText}`; + } + + const collectionWidget = findWidget(node, "collection_name"); + const fieldWidget = findWidget(node, "flattened_field_name"); + const recordIdWidget = findWidget(node, "record_id"); + const indexWidget = findWidget(node, "selected_record_index"); + + if (!collectionWidget || !fieldWidget || !recordIdWidget || !indexWidget) { + console.warn("ZMongoTabularRecordViewNode widgets not found."); + return; + } + + const container = document.createElement("div"); + container.style.display = "flex"; + container.style.flexDirection = "column"; + container.style.gap = "6px"; + container.style.width = "100%"; + container.style.minHeight = "260px"; + container.style.maxHeight = "520px"; + container.style.background = "#111"; + container.style.border = "1px solid #333"; + container.style.borderRadius = "8px"; + container.style.padding = "8px"; + container.style.boxSizing = "border-box"; + + const infoBar = document.createElement("div"); + infoBar.style.fontSize = "12px"; + infoBar.style.color = "#aaa"; + infoBar.textContent = "Loading records..."; + + const searchModeLabel = document.createElement("div"); + searchModeLabel.style.fontSize = "11px"; + searchModeLabel.style.color = "#8fb3ff"; + searchModeLabel.textContent = "Search target: all fields"; + + const searchBar = document.createElement("div"); + searchBar.style.display = "flex"; + searchBar.style.gap = "8px"; + searchBar.style.alignItems = "center"; + searchBar.style.flexWrap = "wrap"; + + const exactSearchInput = document.createElement("input"); + exactSearchInput.type = "text"; + exactSearchInput.placeholder = "Search values (* and ? wildcards supported)"; + exactSearchInput.style.flex = "1"; + exactSearchInput.style.minWidth = "220px"; + exactSearchInput.style.padding = "6px 8px"; + exactSearchInput.style.borderRadius = "6px"; + exactSearchInput.style.border = "1px solid #444"; + exactSearchInput.style.background = "#1b1b1b"; + exactSearchInput.style.color = "#ddd"; + exactSearchInput.style.outline = "none"; + + const exactSearchButton = document.createElement("button"); + exactSearchButton.textContent = "Search"; + exactSearchButton.style.padding = "6px 10px"; + exactSearchButton.style.borderRadius = "6px"; + exactSearchButton.style.border = "1px solid #555"; + exactSearchButton.style.background = "#2a2a2a"; + exactSearchButton.style.color = "#eee"; + exactSearchButton.style.cursor = "pointer"; + + const clearSearchButton = document.createElement("button"); + clearSearchButton.textContent = "Show All"; + clearSearchButton.style.padding = "6px 10px"; + clearSearchButton.style.borderRadius = "6px"; + clearSearchButton.style.border = "1px solid #555"; + clearSearchButton.style.background = "#2a2a2a"; + clearSearchButton.style.color = "#eee"; + clearSearchButton.style.cursor = "pointer"; + + searchBar.appendChild(exactSearchInput); + searchBar.appendChild(exactSearchButton); + searchBar.appendChild(clearSearchButton); + + const tableWrap = document.createElement("div"); + tableWrap.style.flex = "1"; + tableWrap.style.overflow = "auto"; + tableWrap.style.border = "1px solid #333"; + tableWrap.style.borderRadius = "6px"; + tableWrap.style.background = "#181818"; + + const table = document.createElement("table"); + table.style.borderCollapse = "collapse"; + table.style.width = "max-content"; + table.style.minWidth = "100%"; + table.style.tableLayout = "auto"; + table.style.fontSize = "12px"; + + tableWrap.appendChild(table); + container.appendChild(infoBar); + container.appendChild(searchModeLabel); + container.appendChild(searchBar); + container.appendChild(tableWrap); + + if (node.addDOMWidget) { + node.addDOMWidget("record_table", "div", container, { + serialize: false, + hideOnZoom: false, + }); + } + + let lastPayload = { + headings: [], + flat_records: [], + record_ids: [], + default_record_id: "", + record_count: 0, + }; + + function getSafeRecordIds() { + return Array.isArray(lastPayload.record_ids) ? lastPayload.record_ids : []; + } + + function resolveSelectedIndex() { + const recordIds = getSafeRecordIds(); + const recordId = String(recordIdWidget.value || "").trim(); + + if (recordId) { + const idx = recordIds.indexOf(recordId); + if (idx >= 0) { + return idx; + } + } + + return clampIndex(indexWidget.value, recordIds.length); + } + + function syncWidgetsFromIndex(preferredIndex = null) { + const recordIds = getSafeRecordIds(); + + if (!recordIds.length) { + indexWidget.value = 0; + recordIdWidget.value = ""; + node.setDirtyCanvas(true, true); + return 0; + } + + const idx = clampIndex( + preferredIndex !== null ? preferredIndex : indexWidget.value, + recordIds.length + ); + + indexWidget.value = idx; + recordIdWidget.value = recordIds[idx] || ""; + node.setDirtyCanvas(true, true); + return idx; + } + + function renderTable() { + table.innerHTML = ""; + + const headings = Array.isArray(lastPayload.headings) ? lastPayload.headings : []; + const records = Array.isArray(lastPayload.flat_records) ? lastPayload.flat_records : []; + const recordIds = getSafeRecordIds(); + + if (!headings.length || !records.length) { + infoBar.textContent = "No records found."; + return; + } + + const selectedIndex = resolveSelectedIndex(); + infoBar.textContent = `${records.length} record(s) loaded. Click a row to select it. Click a header to set flattened_field_name.`; + + const thead = document.createElement("thead"); + const headRow = document.createElement("tr"); + + const indexHeader = buildCell("#", true); + indexHeader.inner.style.minWidth = "50px"; + indexHeader.inner.style.maxWidth = "80px"; + headRow.appendChild(indexHeader.cell); + + for (const heading of headings) { + const { cell, inner } = buildCell(heading, true); + inner.title = heading; + inner.addEventListener("click", () => { + fieldWidget.value = heading; + updateSearchModeLabel(); + node.setDirtyCanvas(true, true); + }); + headRow.appendChild(cell); + } + + thead.appendChild(headRow); + table.appendChild(thead); + + const tbody = document.createElement("tbody"); + + records.forEach((record, rowIndex) => { + const tr = document.createElement("tr"); + tr.style.cursor = "pointer"; + + if (rowIndex === selectedIndex) { + tr.style.background = "#24364d"; + } + + tr.addEventListener("click", () => { + indexWidget.value = rowIndex; + recordIdWidget.value = recordIds[rowIndex] || ""; + renderTable(); + node.setDirtyCanvas(true, true); + }); + + const indexCell = buildCell(String(rowIndex), false); + indexCell.inner.style.minWidth = "50px"; + indexCell.inner.style.maxWidth = "80px"; + tr.appendChild(indexCell.cell); + + for (const heading of headings) { + let value = record[heading]; + + if (value === null || value === undefined) { + value = ""; + } else if (typeof value !== "string") { + try { + value = JSON.stringify(value, null, 2); + } catch { + value = String(value); + } + } + + const { cell, inner } = buildCell(value, false); + inner.title = value; + tr.appendChild(cell); + } + + tbody.appendChild(tr); + }); + + table.appendChild(tbody); + } + + async function refreshTable() { + const collectionName = collectionWidget.value; + lastPayload = await fetchTabularRecords(collectionName); + + const recordIds = getSafeRecordIds(); + if (recordIds.length) { + const currentId = String(recordIdWidget.value || "").trim(); + + if (!currentId || !recordIds.includes(currentId)) { + recordIdWidget.value = lastPayload.default_record_id || recordIds[0] || ""; + } + + const idx = recordIds.indexOf(recordIdWidget.value); + indexWidget.value = idx >= 0 ? idx : 0; + } else { + recordIdWidget.value = ""; + indexWidget.value = 0; + } + + renderTable(); + node.setDirtyCanvas(true, true); + } + + async function refreshTableSearch(searchText) { + const collectionName = collectionWidget.value; + const flattenedFieldName = String(fieldWidget.value || "").trim(); + + lastPayload = await fetchTabularRecordsSearch( + collectionName, + searchText, + flattenedFieldName + ); + + const recordIds = getSafeRecordIds(); + if (recordIds.length) { + recordIdWidget.value = lastPayload.default_record_id || recordIds[0] || ""; + indexWidget.value = 0; + } else { + recordIdWidget.value = ""; + indexWidget.value = 0; + } + + renderTable(); + node.setDirtyCanvas(true, true); + } + + exactSearchButton.addEventListener("click", async () => { + updateSearchModeLabel(); + try { + await refreshTableSearch(exactSearchInput.value); + updateSearchModeLabel(); + infoBar.textContent = `Search returned ${lastPayload.record_count || 0} record(s).`; + } catch (err) { + console.error("Search failed:", err); + infoBar.textContent = `Error: ${err.message}`; + } + }); + + exactSearchInput.addEventListener("keydown", async (event) => { + if (event.key === "Enter") { + event.preventDefault(); + updateSearchModeLabel(); + try { + await refreshTableSearch(exactSearchInput.value); + updateSearchModeLabel(); + infoBar.textContent = `Search returned ${lastPayload.record_count || 0} record(s).`; + } catch (err) { + console.error("Search failed:", err); + infoBar.textContent = `Error: ${err.message}`; + } + } + }); + + clearSearchButton.addEventListener("click", async () => { + updateSearchModeLabel(); + try { + exactSearchInput.value = ""; + await refreshTable(); + updateSearchModeLabel(); + } catch (err) { + console.error("Refresh all records failed:", err); + infoBar.textContent = `Error: ${err.message}`; + } + }); + + const originalCollectionCallback = collectionWidget.callback; + collectionWidget.callback = async (...args) => { + if (originalCollectionCallback) { + originalCollectionCallback.apply(collectionWidget, args); + } + updateSearchModeLabel(); + try { + exactSearchInput.value = ""; + await refreshTable(); + } catch (err) { + console.error("Failed to refresh tabular records after collection change:", err); + infoBar.textContent = `Error: ${err.message}`; + } + }; + + const originalFieldCallback = fieldWidget.callback; + fieldWidget.callback = (...args) => { + if (originalFieldCallback) { + originalFieldCallback.apply(fieldWidget, args); + } + updateSearchModeLabel(); + node.setDirtyCanvas(true, true); + }; + + const originalIndexCallback = indexWidget.callback; + indexWidget.callback = (...args) => { + if (originalIndexCallback) { + originalIndexCallback.apply(indexWidget, args); + } + syncWidgetsFromIndex(indexWidget.value); + renderTable(); + }; + + const originalRecordIdCallback = recordIdWidget.callback; + recordIdWidget.callback = (...args) => { + if (originalRecordIdCallback) { + originalRecordIdCallback.apply(recordIdWidget, args); + } + + const ids = getSafeRecordIds(); + const idx = ids.indexOf(String(recordIdWidget.value || "").trim()); + + if (idx >= 0) { + indexWidget.value = idx; + } else if (ids.length) { + recordIdWidget.value = ids[0]; + indexWidget.value = 0; + } else { + recordIdWidget.value = ""; + indexWidget.value = 0; + } + + renderTable(); + node.setDirtyCanvas(true, true); + }; + updateSearchModeLabel(); + try { + await refreshTable(); + } catch (err) { + console.error("Initial ZMongo tabular record refresh failed:", err); + infoBar.textContent = `Error: ${err.message}`; + } + + setTimeout(async () => { + updateSearchModeLabel(); + try { + await refreshTable(); + } catch (err) { + console.error("Delayed ZMongo tabular record refresh failed:", err); + } + }, 200); + }, +}); \ No newline at end of file diff --git a/web/js/zmongo_universal_adapter.js b/web/js/zmongo_universal_adapter.js new file mode 100644 index 0000000..bbe6e39 --- /dev/null +++ b/web/js/zmongo_universal_adapter.js @@ -0,0 +1,88 @@ +import { app } from "../../../scripts/app.js"; + +app.registerExtension({ + name: "ZMongo.UniversalModelAdapter", + + async beforeRegisterNodeDef(nodeType, nodeData) { + if (nodeData.name !== "ZMongoUniversalModelAdapterNode") { + return; + } + + const onNodeCreated = nodeType.prototype.onNodeCreated; + nodeType.prototype.onNodeCreated = function () { + const result = onNodeCreated ? onNodeCreated.apply(this, arguments) : undefined; + + this.zmongoActiveOutputs = []; + this.zmongoCompatible = false; + + const statusWidget = this.addWidget( + "text", + "adapter_status", + "No model metadata yet", + () => {}, + { multiline: true } + ); + statusWidget.readOnly = true; + + const updateVisualState = () => { + const active = new Set(this.zmongoActiveOutputs || []); + const outputs = this.outputs || []; + + for (const output of outputs) { + if (!output) continue; + const isActive = active.has(output.name); + + // Visual hint only + output.label = isActive + ? `${output.name} ✓` + : `${output.name} ✗`; + + output.color_on = isActive ? "#6aa84f" : "#666666"; + output.color_off = isActive ? "#38761d" : "#333333"; + } + + this.bgcolor = this.zmongoCompatible ? "#1f3b24" : "#3b1f1f"; + this.setDirtyCanvas(true, true); + }; + + const onExecuted = this.onExecuted; + this.onExecuted = function (message) { + if (onExecuted) { + onExecuted.apply(this, arguments); + } + + try { + // RETURN_NAMES order: + // model, clip, vae, status_text, active_outputs_json, is_compatible, model_type + const statusText = message?.text?.[0] ?? ""; + const activeOutputsJson = message?.text?.[1] ?? ""; + const compatible = message?.text?.[2]; + + const widget = this.widgets?.find((w) => w.name === "adapter_status"); + if (widget) { + widget.value = statusText || "No status"; + } + + let parsed = {}; + try { + parsed = activeOutputsJson ? JSON.parse(activeOutputsJson) : {}; + } catch { + parsed = {}; + } + + this.zmongoActiveOutputs = parsed.active_outputs || []; + this.zmongoCompatible = Boolean( + parsed.compatible ?? compatible ?? false + ); + + updateVisualState(); + } catch (err) { + console.warn("ZMongo adapter UI update failed:", err); + } + }; + + updateVisualState(); + return result; + }; + }, +}); \ No newline at end of file