diff --git a/__init__.py b/__init__.py index 9f38d48..94bebe6 100644 --- a/__init__.py +++ b/__init__.py @@ -1,142 +1,5 @@ -import logging -import os -import sys +from .nodes.zmongo_nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS -from src.zmongo_toolbag.local_vector_search import LocalVectorSearch +WEB_DIRECTORY = "./js" -from src.zmongo_toolbag.zembedder import ZEmbedder - -from src.zmongo_toolbag.data_processor import DataProcessor - -from src.zmongo_toolbag.safe_result import SafeResult - -from src.zmongo_toolbag.zmongo import ZMongo - -# Ensure the custom node package directory is importable so `.src...` works -NODE_PATH = os.path.dirname(__file__) -if NODE_PATH not in sys.path: - sys.path.append(NODE_PATH) - -WEB_DIRECTORY = "./web" - -logger = logging.getLogger(__name__) - - - -from src.zmongo_nodes import ( - ZMongoConfigNode, - ZMongoDataPassThroughNode, - ZMongoDatabaseBrowserNode, - ZMongoFieldSelector, - ZMongoOperationsNode, - ZMongoRecordLoopNode, - ZMongoRecordSplitter, - ZMongoSaveBatchTextNode, - ZMongoSaveTextNode, - ZMongoSaveValueNode, - ZMongoTextFetcher, - ZRetrieverNode, -) -from src.zmongo_field_selector_node import ZMongoFieldSelectorNode -from src.zmongo_flattened_field_selector_node import ZMongoFlattenedFieldSelectorNode -from src.zmongo_tabular_record_view_node import ZMongoTabularRecordViewNode - -from src.zmongo_record_editor_node import ZMongoRecordEditorNode - - -from src.model_loader import ( - ZMongoModelCompatibilityDisplayNode, - ZMongoModelIntrospectorNode, - ZMongoUniversalModelSelectorNode, - ZMongoUniversalModelLoaderNode, - ZMongoUniversalModelAdapterNode, - ZMongoBuiltInLoaderAdapterNode, -) - - - - -############################################################################### -# ComfyUI Required API Registration -############################################################################### -# -# def _safe_register(route_func, label: str) -> None: -# """Register PromptServer routes without breaking node loading.""" -# try: -# route_func(PromptServer.instance) -# except Exception as exc: -# logger.warning("ZMongo %s route registration warning: %s", label, exc) -# -# from .src.zmongo_record_editor_api import register_zmongo_record_editor_routes -# from .src.preset_api import register_preset_routes -# from .src.zmongo_tabular_record_view_api import ( -# register_zmongo_tabular_record_view_routes, -# ) -# from .src.zmongo_field_selector_api import register_zmongo_field_selector_routes -# from server import PromptServer -# -# _safe_register(register_preset_routes, "preset") -# _safe_register(register_zmongo_field_selector_routes, "field selector") -# _safe_register(register_zmongo_tabular_record_view_routes, "tabular record view") -# _safe_register(register_zmongo_record_editor_routes, "record editor") -############################################################################### - -NODE_CLASS_MAPPINGS = { - "ZMongoConfig": ZMongoConfigNode, - "ZMongoTextFetcher": ZMongoTextFetcher, - "ZMongoOperations": ZMongoOperationsNode, - "ZMongoRecordSplitter": ZMongoRecordSplitter, - "ZMongoFieldSelector": ZMongoFieldSelector, - "ZRetriever": ZRetrieverNode, - "ZMongoDatabaseBrowser": ZMongoDatabaseBrowserNode, - "ZMongoRecordLoopNode": ZMongoRecordLoopNode, - "ZMongoSaveTextNode": ZMongoSaveTextNode, - "ZMongoDataPassThroughNode": ZMongoDataPassThroughNode, - "ZMongoSaveBatchTextNode": ZMongoSaveBatchTextNode, - "ZMongoSaveValueNode": ZMongoSaveValueNode, - "ZMongoFlattenedFieldSelectorNode": ZMongoFlattenedFieldSelectorNode, - "ZMongoFieldSelectorNode": ZMongoFieldSelectorNode, - "ZMongoTabularRecordViewNode": ZMongoTabularRecordViewNode, - "ZMongoRecordEditorNode": ZMongoRecordEditorNode, - "ZMongoUniversalModelSelectorNode": ZMongoUniversalModelSelectorNode, - "ZMongoModelIntrospectorNode": ZMongoModelIntrospectorNode, - "ZMongoModelCompatibilityDisplayNode": ZMongoModelCompatibilityDisplayNode, - "ZMongoUniversalModelLoaderNode": ZMongoUniversalModelLoaderNode, - "ZMongoUniversalModelAdapterNode": ZMongoUniversalModelAdapterNode, - "ZMongoBuiltInLoaderAdapterNode": ZMongoBuiltInLoaderAdapterNode, -} - -NODE_DISPLAY_NAME_MAPPINGS = { - "ZMongoConfig": "πŸ›‘οΈ ZMongo Configuration", - "ZMongoTextFetcher": "πŸ“ ZMongo Text Fetcher", - "ZMongoOperations": "πŸ“₯ ZMongo Operations", - "ZMongoRecordSplitter": "βœ‚οΈ ZMongo Record Splitter", - "ZMongoFieldSelector": "πŸ” ZMongo Field Selector", - "ZRetriever": "🧠 ZMongo Vector Retriever", - "ZMongoDatabaseBrowser": "πŸ—‚οΈ ZMongo Database Browser", - "ZMongoRecordLoopNode": "πŸ” ZMongo Record Loop Node", - "ZMongoSaveTextNode": "πŸ’Ύ ZMongo Save Text Node", - "ZMongoDataPassThroughNode": "πŸ”€ ZMongo Data Pass Through Node", - "ZMongoSaveBatchTextNode": "πŸ“š ZMongo Save Batch Text Node", - "ZMongoSaveValueNode": "🧷 ZMongo Save Value Node", - "ZMongoFlattenedFieldSelectorNode": "🧭 ZMongo Flattened Field Selector", - "ZMongoFieldSelectorNode": "πŸ”Ž ZMongo Field Selector Node", - "ZMongoTabularRecordViewNode": "πŸ“Š ZMongo Tabular Record View", - "ZMongoRecordEditorNode": "πŸ“ ZMongo Record Editor", - "ZMongoUniversalModelSelectorNode": "🧩 ZMongo Universal Model Selector", - "ZMongoModelIntrospectorNode": "🧠 ZMongo Model Introspector", - "ZMongoModelCompatibilityDisplayNode": "βœ… ZMongo Model Compatibility Display", - "ZMongoUniversalModelAdapterNode": " ZMongo Universal Model Adapter", - "ZMongoBuiltInLoaderAdapterNode": " ZMongo BuiltIn Loader Adapter", -} - -__all__ = [ - "NODE_CLASS_MAPPINGS", - "NODE_DISPLAY_NAME_MAPPINGS", - "WEB_DIRECTORY", - "DataProcessor", - "SafeResult", - "ZMongo", - "ZEmbedder", - "LocalVectorSearch", -] \ No newline at end of file +__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS', 'WEB_DIRECTORY'] \ No newline at end of file diff --git a/__init__.py.040126_1928 b/__init__.py.040126_1928 new file mode 100644 index 0000000..3001827 --- /dev/null +++ b/__init__.py.040126_1928 @@ -0,0 +1,47 @@ +import os +import sys +import logging + +# 1. PATH CONFIGURATION πŸ› οΈ +# This allows internal files to find 'zmongo_toolbag' and 'gui' easily. +ROOT_PATH = os.path.dirname(__file__) +if ROOT_PATH not in sys.path: + sys.path.append(ROOT_PATH) + +# 2. WEB ASSET REGISTRATION 🌐 +# Tells ComfyUI where to find your JavaScript files for the custom UI. +WEB_DIRECTORY = "./web" + +logger = logging.getLogger(__name__) + +# 3. NODE & API IMPORTS πŸ“¦ +# We import from our sub-packages. These rely on the __init__.py files +# you have inside the 'nodes' and 'gui' folders. +from .nodes import ( + NODE_CLASS_MAPPINGS as BASE_NODES, + NODE_DISPLAY_NAME_MAPPINGS as BASE_DISPLAY, +) + +# register_zmongo_field_selector_routes, +# register_zmongo_record_editor_routes, +# register_zmongo_tabular_record_view_routes, +from server.zmongo_field_selector_api import register_zmongo_field_selector_routes +from server.zmongo_record_editor_api import register_zmongo_record_editor_routes +from server.zmongo_tabular_record_view_api import register_zmongo_tabular_record_view_routes + +# 4. SERVER API REGISTRATION 🚦 +# This connects your Python logic to the ComfyUI web server. +try: + from .server import PromptServer + register_zmongo_field_selector_routes(PromptServer.instance) + register_zmongo_record_editor_routes(PromptServer.instance) + register_zmongo_tabular_record_view_routes(PromptServer.instance) +except Exception as e: + logger.error(f"Failed to register ZMongo API routes: {e}") + +# 5. COMFYUI EXPORTS πŸ“€ +# These are the two variables ComfyUI looks for to load your nodes. +NODE_CLASS_MAPPINGS = {**BASE_NODES} +NODE_DISPLAY_NAME_MAPPINGS = {**BASE_DISPLAY} + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"] \ No newline at end of file diff --git a/api/__init__.py b/api/__init__.py new file mode 100644 index 0000000..ee75ac7 --- /dev/null +++ b/api/__init__.py @@ -0,0 +1,23 @@ +# 1. ROUTE REGISTRATION IMPORTS 🌐 +# We use absolute imports to ensure these work correctly during standalone testing. +from ComfyUI.custom_nodes..api.zmongo_field_selector_api import register_zmongo_field_selector_routes +from api.zmongo_record_editor_api import register_zmongo_record_editor_routes +from api.zmongo_tabular_record_view_api import register_zmongo_tabular_record_view_routes + +# 2. BUNDLED REGISTRATION HELPER 🚦 +# This single function can be called in your root __init__.py to set up everything at once. +def register_all_zmongo_routes(prompt_server_instance): + """ + Registers all ZMongo-related API routes with the ComfyUI server. + """ + register_zmongo_field_selector_routes(prompt_server_instance) + register_zmongo_record_editor_routes(prompt_server_instance) + register_zmongo_tabular_record_view_routes(prompt_server_instance) + +# 3. PUBLIC EXPORTS πŸ“€ +__all__ = [ + "register_all_zmongo_routes", + "register_zmongo_field_selector_routes", + "register_zmongo_record_editor_routes", + "register_zmongo_tabular_record_view_routes", +] \ No newline at end of file diff --git a/api/zmongo_field_selector_api.py b/api/zmongo_field_selector_api.py new file mode 100644 index 0000000..8b46e23 --- /dev/null +++ b/api/zmongo_field_selector_api.py @@ -0,0 +1,37 @@ +import logging +from aiohttp import web + +from zmongo_field_selector_node import ZMongoFieldSelectorNode + +logger = logging.getLogger(__name__) + + +def register_zmongo_field_selector_routes(prompt_server_instance): + if not prompt_server_instance: + logger.warning("PromptServer instance not available; ZMongo routes not registered.") + return + + routes = prompt_server_instance.routes + + @routes.get("/zmongo/flattened_fields") + async def zmongo_get_flattened_fields(request): + try: + collection_name = request.rel_url.query.get("collection_name", "").strip() + fields = ZMongoFieldSelectorNode.get_flattened_field_names(collection_name) + return web.json_response( + { + "success": True, + "collection_name": collection_name, + "fields": fields, + } + ) + except Exception as exc: + logger.exception("Failed to get flattened fields: %s", exc) + return web.json_response( + { + "success": False, + "error": str(exc), + "fields": [], + }, + status=500, + ) \ No newline at end of file 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.json b/example_workflows/zmongo_db_browser.json new file mode 100644 index 0000000..85a2aeb --- /dev/null +++ b/example_workflows/zmongo_db_browser.json @@ -0,0 +1 @@ +{"id":"9be9b786-3a4f-4ed3-b0f8-c81334a58dd9","revision":0,"last_node_id":12,"last_link_id":11,"nodes":[{"id":5,"type":"ZMongoDatabaseBrowser","pos":[-1587.0391233787002,1217.7367585716897],"size":[301.71132507324216,278],"flags":{},"order":0,"mode":0,"inputs":[{"localized_name":"mongo_uri","name":"mongo_uri","type":"STRING","widget":{"name":"mongo_uri"},"link":null},{"localized_name":"database_name","name":"database_name","type":"STRING","widget":{"name":"database_name"},"link":null},{"localized_name":"collection_name","name":"collection_name","type":"STRING","widget":{"name":"collection_name"},"link":null},{"localized_name":"limit","name":"limit","type":"INT","widget":{"name":"limit"},"link":null},{"localized_name":"select_index","name":"select_index","type":"INT","widget":{"name":"select_index"},"link":null},{"localized_name":"refresh_nonce","name":"refresh_nonce","type":"INT","widget":{"name":"refresh_nonce"},"link":null}],"outputs":[{"localized_name":"selected_record_json","name":"selected_record_json","type":"STRING","links":[9]},{"localized_name":"record_list_summary","name":"record_list_summary","type":"STRING","links":[8]},{"localized_name":"total_count","name":"total_count","type":"INT","links":[7]},{"localized_name":"db_name_out","name":"db_name_out","type":"STRING","links":[10]},{"localized_name":"coll_name_out","name":"coll_name_out","type":"STRING","links":[6]},{"localized_name":"selected_index_out","name":"selected_index_out","type":"INT","links":[11]}],"properties":{"aux_id":"CentralFloridaAttorney/ComfyUI-ZMongo","ver":"ce4a054405011239823219cdc868a20aafc21214","Node name for S&R":"ZMongoDatabaseBrowser"},"widgets_values":["mongodb://127.0.0.1:27017","test","ocr_docs",50,3,1]},{"id":9,"type":"PreviewAny","pos":[-1002.9658276151989,295.91699285067153],"size":[376.7605899798724,590.0483733678761],"flags":{},"order":2,"mode":0,"inputs":[{"localized_name":"source","name":"source","type":"*","link":9}],"outputs":[],"title":"Selected Record","properties":{"cnr_id":"comfy-core","ver":"0.18.1","Node name for S&R":"PreviewAny"},"widgets_values":[null,null,null]},{"id":6,"type":"PreviewAny","pos":[-991.8215459206516,939.8597664149295],"size":[210,112],"flags":{},"order":4,"mode":0,"inputs":[{"localized_name":"source","name":"source","type":"*","link":7}],"outputs":[],"title":"Total Records","properties":{"cnr_id":"comfy-core","ver":"0.18.1","Node name for S&R":"PreviewAny"},"widgets_values":[null,null,null]},{"id":11,"type":"PreviewAny","pos":[-1085.4523439654295,1875.885513541283],"size":[210,166],"flags":{},"order":7,"mode":0,"inputs":[{"localized_name":"source","name":"source","type":"*","link":11}],"outputs":[],"title":"Selected Index","properties":{"cnr_id":"comfy-core","ver":"0.18.1","Node name for S&R":"PreviewAny"},"widgets_values":[null,null,null]},{"id":7,"type":"PreviewAny","pos":[-1073.064410849724,1569.9987484820574],"size":[210,166],"flags":{},"order":6,"mode":0,"inputs":[{"localized_name":"source","name":"source","type":"*","link":6}],"outputs":[],"title":"Collection Name","properties":{"cnr_id":"comfy-core","ver":"0.18.1","Node name for S&R":"PreviewAny"},"widgets_values":[null,null,null]},{"id":10,"type":"PreviewAny","pos":[-1135.9570247493903,1290.7940096692307],"size":[210,112],"flags":{},"order":5,"mode":0,"inputs":[{"localized_name":"source","name":"source","type":"*","link":10}],"outputs":[],"title":"Database Name","properties":{"cnr_id":"comfy-core","ver":"0.18.1","Node name for S&R":"PreviewAny"},"widgets_values":[null,null,null]},{"id":8,"type":"PreviewAny","pos":[-578.5437776325792,985.2210211795292],"size":[443.4647667129609,497.6153491606501],"flags":{},"order":3,"mode":0,"inputs":[{"localized_name":"source","name":"source","type":"*","link":8}],"outputs":[],"title":"Records Preview","properties":{"cnr_id":"comfy-core","ver":"0.18.1","Node name for S&R":"PreviewAny"},"widgets_values":[null,null,null]},{"id":12,"type":"MarkdownNote","pos":[-2487.135818122622,470.2028312144485],"size":[1203.11704502172,644.5748535982483],"flags":{},"order":1,"mode":0,"inputs":[],"outputs":[],"title":"Install MongoDB","properties":{},"widgets_values":["To set up the backend for your database nodes (like **ZMongo**), you'll need the **MongoDB Community Edition**. As of 2026, the current stable version is **8.x**, which includes native vector search supportβ€”perfect for AI-driven projects.\n\n### 🌐 Official Download Portal\nThe easiest way to get the correct version for any system is the **MongoDB Download Center**:\nπŸ‘‰ **[Download MongoDB Community Server](https://www.mongodb.com/try/download/community)**\n\n---\n\n### πŸ’» Platform-Specific Instructions\n\n#### **Windows**\nFor Windows 10 or 11, it is recommended to use the **MSI Installer**.\n1. Select **Windows** as the platform on the download page.\n2. Download and run the `.msi` file.\n3. **Pro Tip:** During installation, ensure \"**Install MongoDB as a Service**\" is checked so it starts automatically with your computer.\n\n#### **macOS**\nThe standard way to install on Mac is via **Homebrew**. Open your terminal and run:\n```bash\nbrew tap mongodb/brew\nbrew install mongodb-community@8.0\n```\nTo start the database: `brew services start mongodb-community@8.0`.\n\n#### **Linux**\nMongoDB provides native packages (`.deb` and `.rpm`) for major distributions like **Ubuntu, Debian, and RedHat**.\n* **Ubuntu/Debian:** Use the `.deb` package.\n* **RedHat/CentOS:** Use the `.rpm` package.\n* Detailed guides for each distro can be found in the [Official Linux Installation Docs](https://www.mongodb.com/docs/manual/administration/install-on-linux/).\n\n#### **Docker (Quickest Setup)**\nIf you don't want to install it directly on your OS, you can run a containerized version with one command:\n```bash\ndocker pull mongodb/mongodb-community-server\ndocker run --name mongodb -d -p 27017:27017 mongodb/mongodb-community-server\n```\n\n---\n\n### πŸ” Don't Forget the GUI\nWhile the server runs in the background, you'll likely want **MongoDB Compass**. It’s the official visual tool for exploring your data, and it is usually offered as an optional checkbox during the Windows installation.\n\n* **Download Compass:** [https://www.mongodb.com/try/download/compass](https://www.mongodb.com/try/download/compass)\n\n> [!NOTE]\n> Since you are using **ComfyUI-ZMongo**, ensure your connection string in ComfyUI matches your local instance (usually `mongodb://127.0.0.1:27017`) once the service is running."],"color":"#432","bgcolor":"#653"}],"links":[[6,5,4,7,0,"STRING"],[7,5,2,6,0,"INT"],[8,5,1,8,0,"STRING"],[9,5,0,9,0,"STRING"],[10,5,3,10,0,"STRING"],[11,5,5,11,0,"INT"]],"groups":[],"config":{},"extra":{"ds":{"scale":0.8427069746815818,"offset":[2950.578449003553,-305.10825612995717]}},"version":0.4} \ No newline at end of file diff --git a/example_workflows/zmongo_db_browser.png b/example_workflows/zmongo_db_browser.png new file mode 100644 index 0000000..6c37873 Binary files /dev/null and b/example_workflows/zmongo_db_browser.png differ diff --git a/example_workflows/zmongo_loop_sd15_demo.json b/example_workflows/zmongo_loop_sd15_demo.json new file mode 100644 index 0000000..d997622 --- /dev/null +++ b/example_workflows/zmongo_loop_sd15_demo.json @@ -0,0 +1,254 @@ +{ + "last_node_id": 9, + "last_link_id": 9, + "nodes": [ + { + "id": 1, + "type": "ZMongoRecordLoopNode", + "pos": [-1200, 200], + "size": [340, 420], + "flags": {}, + "order": 0, + "mode": 0, + "inputs": [], + "outputs": [ + { "name": "record_data", "type": "STRING", "links": [] }, + { "name": "selected_field_text", "type": "STRING", "links": [1] }, + { "name": "current_index", "type": "INT", "links": [] }, + { "name": "total_count", "type": "INT", "links": [] }, + { "name": "emitted_count", "type": "INT", "links": [] }, + { "name": "summary", "type": "STRING", "links": [8] }, + { "name": "loop_state_json", "type": "STRING", "links": [] } + ], + "properties": {}, + "widgets_values": [ + "mongodb://127.0.0.1:27017", + "test", + "ocr_docs", + "range", + 1, + 10, + 1, + 1, + "_id", + "ascending", + "text", + "text\ndoc_text\ndocument.text\nprompt", + "{}", + "{}", + false, + 0, + "fixed", + true, + true, + true, + 500, + "use_placeholder", + "prompt_text", + "wrap", + "record_json", + 0 + ], + "_meta": { + "title": "πŸ” ZMongo Record Loop" + } + }, + { + "id": 2, + "type": "CheckpointLoaderSimple", + "pos": [-700, 200], + "size": [300, 120], + "flags": {}, + "order": 1, + "mode": 0, + "inputs": [], + "outputs": [ + { "name": "MODEL", "type": "MODEL", "links": [3] }, + { "name": "CLIP", "type": "CLIP", "links": [1, 2] }, + { "name": "VAE", "type": "VAE", "links": [6] } + ], + "properties": {}, + "widgets_values": [ + "v1-5-pruned-emaonly-fp16.safetensors" + ], + "_meta": { + "title": "SD1.5 Checkpoint" + } + }, + { + "id": 3, + "type": "CLIPTextEncode", + "pos": [-400, 100], + "size": [300, 150], + "flags": {}, + "order": 2, + "mode": 0, + "inputs": [ + { "name": "text", "type": "STRING", "link": 1 }, + { "name": "clip", "type": "CLIP", "link": 2 } + ], + "outputs": [ + { "name": "CONDITIONING", "type": "CONDITIONING", "links": [4] } + ], + "properties": {}, + "_meta": { + "title": "Positive Prompt" + } + }, + { + "id": 4, + "type": "CLIPTextEncode", + "pos": [-400, 300], + "size": [300, 150], + "flags": {}, + "order": 3, + "mode": 0, + "inputs": [ + { "name": "text", "type": "STRING", "widget": { "name": "text" }, "link": null }, + { "name": "clip", "type": "CLIP", "link": 2 } + ], + "outputs": [ + { "name": "CONDITIONING", "type": "CONDITIONING", "links": [5] } + ], + "properties": {}, + "widgets_values": [ + "bad anatomy, blurry, distorted, extra limbs" + ], + "_meta": { + "title": "Negative Prompt" + } + }, + { + "id": 5, + "type": "EmptyLatentImage", + "pos": [-100, 200], + "size": [300, 120], + "flags": {}, + "order": 4, + "mode": 0, + "inputs": [], + "outputs": [ + { "name": "LATENT", "type": "LATENT", "links": [6] } + ], + "properties": {}, + "widgets_values": [ + 512, + 512, + 1 + ], + "_meta": { + "title": "Empty Latent Image" + } + }, + { + "id": 6, + "type": "KSampler", + "pos": [200, 200], + "size": [350, 250], + "flags": {}, + "order": 5, + "mode": 0, + "inputs": [ + { "name": "model", "type": "MODEL", "link": 3 }, + { "name": "positive", "type": "CONDITIONING", "link": 4 }, + { "name": "negative", "type": "CONDITIONING", "link": 5 }, + { "name": "latent_image", "type": "LATENT", "link": 6 } + ], + "outputs": [ + { "name": "LATENT", "type": "LATENT", "links": [7] } + ], + "properties": {}, + "widgets_values": [ + 1, + 20, + 7, + "euler", + "normal", + 1 + ], + "_meta": { + "title": "KSampler" + } + }, + { + "id": 7, + "type": "VAEDecode", + "pos": [550, 200], + "size": [300, 120], + "flags": {}, + "order": 6, + "mode": 0, + "inputs": [ + { "name": "samples", "type": "LATENT", "link": 7 }, + { "name": "vae", "type": "VAE", "link": 6 } + ], + "outputs": [ + { "name": "IMAGE", "type": "IMAGE", "links": [9] } + ], + "properties": {}, + "_meta": { + "title": "VAE Decode" + } + }, + { + "id": 8, + "type": "SaveImage", + "pos": [850, 200], + "size": [250, 120], + "flags": {}, + "order": 7, + "mode": 0, + "inputs": [ + { "name": "images", "type": "IMAGE", "link": 9 } + ], + "outputs": [], + "properties": {}, + "widgets_values": [ + "zmongo_loop_output" + ], + "_meta": { + "title": "Save Image" + } + }, + { + "id": 9, + "type": "PreviewAny", + "pos": [-800, 650], + "size": [500, 220], + "flags": {}, + "order": 8, + "mode": 0, + "inputs": [ + { "name": "source", "type": "*", "link": 8 } + ], + "outputs": [], + "properties": {}, + "widgets_values": [ + null, + null, + null + ], + "_meta": { + "title": "Loop Summary" + } + } + ], + "links": [ + [1, 1, 1, 3, 0, "STRING"], + [2, 2, 1, 3, 1, "CLIP"], + [3, 2, 0, 6, 0, "MODEL"], + [4, 3, 0, 6, 1, "CONDITIONING"], + [5, 4, 0, 6, 2, "CONDITIONING"], + [6, 5, 0, 6, 3, "LATENT"], + [7, 6, 0, 7, 0, "LATENT"], + [8, 1, 5, 9, 0, "STRING"], + [9, 7, 0, 8, 0, "IMAGE"] + ], + "groups": [], + "config": {}, + "extra": {}, + "version": 0.4, + "_meta": { + "title": "πŸ” ZMongo β†’ SD1.5 Loop Demo" + } +} \ No newline at end of file diff --git a/example_workflows/zmongo_loop_sd15_demo.png b/example_workflows/zmongo_loop_sd15_demo.png new file mode 100644 index 0000000..6c37873 Binary files /dev/null and b/example_workflows/zmongo_loop_sd15_demo.png differ diff --git a/example_workflows/zmongo_record_loop_demoa - Copy.jpg b/example_workflows/zmongo_record_loop_demoa - Copy.jpg new file mode 100644 index 0000000..6c37873 Binary files /dev/null and b/example_workflows/zmongo_record_loop_demoa - Copy.jpg differ diff --git a/example_workflows/zmongo_record_loop_demoa - Copy.json b/example_workflows/zmongo_record_loop_demoa - Copy.json new file mode 100644 index 0000000..50e5c54 --- /dev/null +++ b/example_workflows/zmongo_record_loop_demoa - Copy.json @@ -0,0 +1,96 @@ +{ + "last_node_id": 3, + "last_link_id": 2, + "nodes": [ + { + "id": 1, + "type": "ZMongoRecordLoopNode", + "pos": [-520, 180], + "size": [380, 460], + "flags": {}, + "order": 0, + "mode": 0, + "inputs": [], + "outputs": [ + { "name": "selected_record_json_or_list", "type": "STRING", "links": [] }, + { "name": "selected_field_value_or_list", "type": "STRING", "links": [1] }, + { "name": "current_record_index", "type": "INT", "links": [2] }, + { "name": "total_record_count", "type": "INT", "links": [] }, + { "name": "returned_item_count", "type": "INT", "links": [] }, + { "name": "summary", "type": "STRING", "links": [] }, + { "name": "selection_state_json", "type": "STRING", "links": [] } + ], + "properties": {}, + "widgets_values": [ + "mongodb://127.0.0.1:27017", + "test", + "ocr_docs", + "record range", + 2, + 1000, + 1, + 1, + "_id", + "ascending", + "text", + "text\ndoc_text\ndocument.text\nprompt", + "{}", + "{}", + false, + 0, + true, + true, + true, + 2000, + "skip record", + "empty prompt", + "stop at end", + "field value list as json", + 0 + ], + "_meta": { + "title": "πŸ” ZMongo Record Selector" + } + }, + { + "id": 2, + "type": "ZMongoSaveTextNode", + "pos": [20, 200], + "size": [360, 240], + "flags": {}, + "order": 1, + "mode": 0, + "inputs": [ + { "name": "content_to_save", "type": "STRING", "link": 1 }, + { "name": "starting_index", "type": "INT", "link": 2 } + ], + "outputs": [ + { "name": "saved_file_path_or_paths", "type": "STRING", "links": [] }, + { "name": "saved_file_count", "type": "INT", "links": [] } + ], + "properties": {}, + "widgets_values": [ + "zmongo_record", + 1, + "txt", + "write new content", + "separate file per item", + true + ], + "_meta": { + "title": "πŸ’Ύ ZMongo Save Text" + } + } + ], + "links": [ + [1, 1, 1, 2, 0, "STRING"], + [2, 1, 2, 2, 1, "INT"] + ], + "groups": [], + "config": {}, + "extra": {}, + "version": 0.4, + "_meta": { + "title": "ZMongo β†’ Save Range of Text Fields" + } +} \ No newline at end of file diff --git a/example_workflows/zmongo_record_loop_demoa.json b/example_workflows/zmongo_record_loop_demoa.json new file mode 100644 index 0000000..50e5c54 --- /dev/null +++ b/example_workflows/zmongo_record_loop_demoa.json @@ -0,0 +1,96 @@ +{ + "last_node_id": 3, + "last_link_id": 2, + "nodes": [ + { + "id": 1, + "type": "ZMongoRecordLoopNode", + "pos": [-520, 180], + "size": [380, 460], + "flags": {}, + "order": 0, + "mode": 0, + "inputs": [], + "outputs": [ + { "name": "selected_record_json_or_list", "type": "STRING", "links": [] }, + { "name": "selected_field_value_or_list", "type": "STRING", "links": [1] }, + { "name": "current_record_index", "type": "INT", "links": [2] }, + { "name": "total_record_count", "type": "INT", "links": [] }, + { "name": "returned_item_count", "type": "INT", "links": [] }, + { "name": "summary", "type": "STRING", "links": [] }, + { "name": "selection_state_json", "type": "STRING", "links": [] } + ], + "properties": {}, + "widgets_values": [ + "mongodb://127.0.0.1:27017", + "test", + "ocr_docs", + "record range", + 2, + 1000, + 1, + 1, + "_id", + "ascending", + "text", + "text\ndoc_text\ndocument.text\nprompt", + "{}", + "{}", + false, + 0, + true, + true, + true, + 2000, + "skip record", + "empty prompt", + "stop at end", + "field value list as json", + 0 + ], + "_meta": { + "title": "πŸ” ZMongo Record Selector" + } + }, + { + "id": 2, + "type": "ZMongoSaveTextNode", + "pos": [20, 200], + "size": [360, 240], + "flags": {}, + "order": 1, + "mode": 0, + "inputs": [ + { "name": "content_to_save", "type": "STRING", "link": 1 }, + { "name": "starting_index", "type": "INT", "link": 2 } + ], + "outputs": [ + { "name": "saved_file_path_or_paths", "type": "STRING", "links": [] }, + { "name": "saved_file_count", "type": "INT", "links": [] } + ], + "properties": {}, + "widgets_values": [ + "zmongo_record", + 1, + "txt", + "write new content", + "separate file per item", + true + ], + "_meta": { + "title": "πŸ’Ύ ZMongo Save Text" + } + } + ], + "links": [ + [1, 1, 1, 2, 0, "STRING"], + [2, 1, 2, 2, 1, "INT"] + ], + "groups": [], + "config": {}, + "extra": {}, + "version": 0.4, + "_meta": { + "title": "ZMongo β†’ Save Range of Text Fields" + } +} \ No newline at end of file diff --git a/example_workflows/zmongo_record_loop_demoa.png b/example_workflows/zmongo_record_loop_demoa.png new file mode 100644 index 0000000..6c37873 Binary files /dev/null and b/example_workflows/zmongo_record_loop_demoa.png differ diff --git a/example_workflows/zmongo_system_demoa.json b/example_workflows/zmongo_system_demoa.json new file mode 100644 index 0000000..92d3a97 --- /dev/null +++ b/example_workflows/zmongo_system_demoa.json @@ -0,0 +1,544 @@ +{ + "last_node_id": 16, + "last_link_id": 14, + "nodes": [ + { + "id": 1, + "type": "MarkdownNote", + "pos": [-2520, -180], + "size": [960, 420], + "flags": {}, + "order": 0, + "mode": 0, + "inputs": [], + "outputs": [], + "title": "ZMongo Workflow Guide", + "properties": {}, + "widgets_values": [ + "# ZMongo Workflow Guide\n\nThis workflow demonstrates the recommended way to use the current ZMongo nodes inside ComfyUI.\n\n## What this template shows\n1. Configure the shared ZMongo connection.\n2. Browse records from a collection.\n3. Inspect the selected record JSON.\n4. Extract a field using the field selector.\n5. Select a record range using the record loop node.\n6. Save the selected field values to files.\n\n## Recommended usage pattern\n- Use **ZMongoConfigNode** first to confirm your database connection.\n- Use **ZMongoDatabaseBrowserNode** when you want to inspect real records and discover fields.\n- Use **ZMongoFieldSelector** to pick a field from a selected record.\n- Use **ZMongoRecordLoopNode** when you want one record, a range of records, or a field-value list for batch export.\n- Use **ZMongoSaveTextNode** to save one value or a JSON list of values to the ComfyUI output folder.\n\n## Typical field names\nFor OCR or document workflows, the most common field is usually:\n- `text`\n\n## Important note about ranges\nThe record loop node returns data according to its `emit_mode`. If you want to save multiple field values at once, use:\n- `emit_mode = list_of_prompts`\n\nThe save node can then split that JSON list into separate files when:\n- `save_layout = separate file per item`\n\n## Expected output folder\nSaved files will go into the ComfyUI output directory." + ], + "color": "#2a2f38", + "bgcolor": "#3a4250" + }, + { + "id": 2, + "type": "ZMongoConfigNode", + "pos": [-2450, 360], + "size": [320, 130], + "flags": {}, + "order": 1, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "status", + "type": "STRING", + "links": [1] + } + ], + "properties": {}, + "widgets_values": [ + "mongodb://127.0.0.1:27017", + "test" + ], + "_meta": { + "title": "ZMongo Config" + } + }, + { + "id": 3, + "type": "PreviewAny", + "pos": [-2070, 350], + "size": [340, 150], + "flags": {}, + "order": 2, + "mode": 0, + "inputs": [ + { + "name": "source", + "type": "*", + "link": 1 + } + ], + "outputs": [], + "properties": {}, + "widgets_values": [null, null, false], + "_meta": { + "title": "Connection Status" + } + }, + { + "id": 4, + "type": "MarkdownNote", + "pos": [-1510, -180], + "size": [760, 360], + "flags": {}, + "order": 3, + "mode": 0, + "inputs": [], + "outputs": [], + "title": "How To Browse And Discover Fields", + "properties": {}, + "widgets_values": [ + "## Browsing workflow\n\nUse the database browser to inspect real records from the collection before building downstream logic.\n\n### Key controls\n- **mongo_uri**: database URI\n- **database_name**: database to inspect\n- **collection_name**: collection to inspect\n- **limit**: how many records to preview\n- **select_index**: which record to fetch as the selected record\n\n### Recommended first step\nStart with:\n- collection: `ocr_docs`\n- limit: `20`\n- select_index: `1`\n\nThen inspect:\n- `selected_record_json`\n- `record_list_summary`\n- `selected_index_out`\n\nAfter that, connect `selected_record_json` into **ZMongoFieldSelector** and choose a field like `text`." + ], + "color": "#2f3a2a", + "bgcolor": "#45563d" + }, + { + "id": 5, + "type": "ZMongoDatabaseBrowserNode", + "pos": [-1490, 240], + "size": [320, 280], + "flags": {}, + "order": 4, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "selected_record_json", + "type": "STRING", + "links": [2, 4] + }, + { + "name": "record_list_summary", + "type": "STRING", + "links": [3] + }, + { + "name": "total_count", + "type": "INT", + "links": [] + }, + { + "name": "db_name_out", + "type": "STRING", + "links": [] + }, + { + "name": "coll_name_out", + "type": "STRING", + "links": [] + }, + { + "name": "selected_index_out", + "type": "INT", + "links": [] + } + ], + "properties": {}, + "widgets_values": [ + "mongodb://127.0.0.1:27017", + "test", + "ocr_docs", + 20, + 1, + 1 + ], + "_meta": { + "title": "ZMongo Database Browser" + } + }, + { + "id": 6, + "type": "PreviewAny", + "pos": [-1080, 220], + "size": [390, 420], + "flags": {}, + "order": 5, + "mode": 0, + "inputs": [ + { + "name": "source", + "type": "*", + "link": 2 + } + ], + "outputs": [], + "properties": {}, + "widgets_values": [null, null, false], + "_meta": { + "title": "Selected Record JSON" + } + }, + { + "id": 7, + "type": "PreviewAny", + "pos": [-1080, 670], + "size": [390, 260], + "flags": {}, + "order": 6, + "mode": 0, + "inputs": [ + { + "name": "source", + "type": "*", + "link": 3 + } + ], + "outputs": [], + "properties": {}, + "widgets_values": [null, null, false], + "_meta": { + "title": "Record List Summary" + } + }, + { + "id": 8, + "type": "ZMongoFieldSelector", + "pos": [-630, 320], + "size": [300, 90], + "flags": {}, + "order": 7, + "mode": 0, + "inputs": [ + { + "name": "record_json", + "type": "STRING", + "link": 4 + } + ], + "outputs": [ + { + "name": "field_value", + "type": "STRING", + "links": [5] + }, + { + "name": "available_paths", + "type": "STRING", + "links": [6] + }, + { + "name": "selected_path", + "type": "STRING", + "links": [7] + } + ], + "properties": {}, + "widgets_values": [ + "text" + ], + "_meta": { + "title": "ZMongo Field Selector" + } + }, + { + "id": 9, + "type": "PreviewAny", + "pos": [-240, 250], + "size": [360, 200], + "flags": {}, + "order": 8, + "mode": 0, + "inputs": [ + { + "name": "source", + "type": "*", + "link": 5 + } + ], + "outputs": [], + "properties": {}, + "widgets_values": [null, null, false], + "_meta": { + "title": "Selected Field Value" + } + }, + { + "id": 10, + "type": "PreviewAny", + "pos": [-240, 480], + "size": [360, 260], + "flags": {}, + "order": 9, + "mode": 0, + "inputs": [ + { + "name": "source", + "type": "*", + "link": 6 + } + ], + "outputs": [], + "properties": {}, + "widgets_values": [null, null, false], + "_meta": { + "title": "Available Field Paths" + } + }, + { + "id": 11, + "type": "PreviewAny", + "pos": [-240, 770], + "size": [240, 120], + "flags": {}, + "order": 10, + "mode": 0, + "inputs": [ + { + "name": "source", + "type": "*", + "link": 7 + } + ], + "outputs": [], + "properties": {}, + "widgets_values": [null, null, false], + "_meta": { + "title": "Selected Path" + } + }, + { + "id": 12, + "type": "MarkdownNote", + "pos": [180, -180], + "size": [900, 420], + "flags": {}, + "order": 11, + "mode": 0, + "inputs": [], + "outputs": [], + "title": "How To Use The Record Loop Node", + "properties": {}, + "widgets_values": [ + "## Record loop workflow\n\nUse the record loop node when you want to select one record or many records for export or downstream generation.\n\n### Recommended meanings of the current fields\n- **mode**\n - `single` = one record\n - `range` = a record range\n - `all` = all matching records\n - `batch` = first batch from the selected range\n - `random_single` = one random record\n - `random_batch` = random batch of records\n\n- **field_path** = main field to extract, usually `text`\n- **fallback_field_paths** = backup field names\n- **emit_mode**\n - `prompt_text` = one field value\n - `record_json` = one record as JSON\n - `record_and_prompt` = one record plus one field value\n - `list_of_records` = JSON list of records\n - `list_of_prompts` = JSON list of field values\n\n### Recommended batch export setting\nTo save a range of `text` values to files, use:\n- `mode = range`\n- `start_index = 1`\n- `end_index = 4`\n- `field_path = text`\n- `emit_mode = list_of_prompts`\n\nThen connect the second output into **ZMongoSaveTextNode**." + ], + "color": "#3a2f2a", + "bgcolor": "#5b463f" + }, + { + "id": 13, + "type": "ZMongoRecordLoopNode", + "pos": [180, 280], + "size": [360, 440], + "flags": {}, + "order": 12, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "record_data", + "type": "STRING", + "links": [] + }, + { + "name": "selected_field_text", + "type": "STRING", + "links": [8] + }, + { + "name": "current_index", + "type": "INT", + "links": [9] + }, + { + "name": "total_count", + "type": "INT", + "links": [] + }, + { + "name": "emitted_count", + "type": "INT", + "links": [] + }, + { + "name": "summary", + "type": "STRING", + "links": [10] + }, + { + "name": "loop_state_json", + "type": "STRING", + "links": [11] + } + ], + "properties": {}, + "widgets_values": [ + "mongodb://127.0.0.1:27017", + "test", + "ocr_docs", + "range", + 1, + 4, + 1, + 1, + "_id", + "ascending", + "text", + "text\ndoc_text\ndocument.text\nprompt", + "{}", + "{}", + false, + 0, + true, + true, + true, + 2000, + "skip", + "empty prompt", + "stop_at_end", + "list_of_prompts", + 0 + ], + "_meta": { + "title": "ZMongo Record Loop" + } + }, + { + "id": 14, + "type": "ZMongoSaveTextNode", + "pos": [650, 320], + "size": [360, 250], + "flags": {}, + "order": 13, + "mode": 0, + "inputs": [ + { + "name": "content_to_save", + "type": "STRING", + "link": 8 + }, + { + "name": "starting_index", + "type": "INT", + "link": 9 + } + ], + "outputs": [ + { + "name": "saved_file_path_or_paths", + "type": "STRING", + "links": [12] + }, + { + "name": "saved_file_count", + "type": "INT", + "links": [13] + } + ], + "properties": {}, + "widgets_values": [ + "zmongo_export", + 1, + "txt", + "write new content", + "separate file per item", + true + ], + "_meta": { + "title": "ZMongo Save Text" + } + }, + { + "id": 15, + "type": "PreviewAny", + "pos": [1050, 260], + "size": [400, 220], + "flags": {}, + "order": 14, + "mode": 0, + "inputs": [ + { + "name": "source", + "type": "*", + "link": 12 + } + ], + "outputs": [], + "properties": {}, + "widgets_values": [null, null, false], + "_meta": { + "title": "Saved File Paths" + } + }, + { + "id": 16, + "type": "PreviewAny", + "pos": [1050, 520], + "size": [300, 140], + "flags": {}, + "order": 15, + "mode": 0, + "inputs": [ + { + "name": "source", + "type": "*", + "link": 13 + } + ], + "outputs": [], + "properties": {}, + "widgets_values": [null, null, false], + "_meta": { + "title": "Saved File Count" + } + }, + { + "id": 17, + "type": "PreviewAny", + "pos": [650, 620], + "size": [360, 220], + "flags": {}, + "order": 16, + "mode": 0, + "inputs": [ + { + "name": "source", + "type": "*", + "link": 10 + } + ], + "outputs": [], + "properties": {}, + "widgets_values": [null, null, false], + "_meta": { + "title": "Loop Summary" + } + }, + { + "id": 18, + "type": "PreviewAny", + "pos": [650, 870], + "size": [360, 260], + "flags": {}, + "order": 17, + "mode": 0, + "inputs": [ + { + "name": "source", + "type": "*", + "link": 11 + } + ], + "outputs": [], + "properties": {}, + "widgets_values": [null, null, false], + "_meta": { + "title": "Loop State JSON" + } + } + ], + "links": [ + [1, 2, 0, 3, 0, "STRING"], + [2, 5, 0, 6, 0, "STRING"], + [3, 5, 1, 7, 0, "STRING"], + [4, 5, 0, 8, 0, "STRING"], + [5, 8, 0, 9, 0, "STRING"], + [6, 8, 1, 10, 0, "STRING"], + [7, 8, 2, 11, 0, "STRING"], + [8, 13, 1, 14, 0, "STRING"], + [9, 13, 2, 14, 1, "INT"], + [10, 13, 5, 17, 0, "STRING"], + [11, 13, 6, 18, 0, "STRING"], + [12, 14, 0, 15, 0, "STRING"], + [13, 14, 1, 16, 0, "INT"] + ], + "groups": [], + "config": {}, + "extra": {}, + "version": 0.4, + "_meta": { + "title": "ZMongo Guided Workflow Template" + } +} \ No newline at end of file diff --git a/example_workflows/zmongo_system_demoa.png b/example_workflows/zmongo_system_demoa.png new file mode 100644 index 0000000..6c37873 Binary files /dev/null and b/example_workflows/zmongo_system_demoa.png differ diff --git a/gui/__init__.py b/gui/__init__.py new file mode 100644 index 0000000..854bdb2 --- /dev/null +++ b/gui/__init__.py @@ -0,0 +1,81 @@ +from .zmongo_toolbag.data_processor import DataProcessor +from .zmongo_toolbag.local_vector_search import LocalVectorSearch +from .zmongo_toolbag.safe_result import SafeResult +from .zmongo_toolbag.zembedder import ZEmbedder +from .zmongo_toolbag.zmongo import ZMongo +from .preset_api import register_preset_routes + +from .zmongo_field_selector_api import register_zmongo_field_selector_routes +from .zmongo_field_selector_node import ZMongoFieldSelectorNode +from .zmongo_flattened_field_selector_node import ZMongoFlattenedFieldSelectorNode +from .zmongo_record_editor_api import register_zmongo_record_editor_routes +from .zmongo_record_editor_node import ZMongoRecordEditorNode +from .zmongo_tabular_record_view_api import register_zmongo_tabular_record_view_routes +from .zmongo_tabular_record_view_node import ZMongoTabularRecordViewNode +from .zmongo_nodes import ( + ZMongoConfigNode, + ZMongoDataPassThroughNode, + ZMongoDatabaseBrowserNode, + ZMongoFieldSelector, + ZMongoOperationsNode, + ZMongoRecordLoopNode, + ZMongoRecordSplitter, + ZMongoSaveBatchTextNode, + ZMongoSaveTextNode, + ZMongoSaveValueNode, + ZMongoTextFetcher, + ZRetrieverNode, +) + +from .model_loader.zmongo_model_compatibility_display_node import ( + ZMongoModelCompatibilityDisplayNode, +) +from .model_loader.zmongo_model_introspector_node import ( + ZMongoModelIntrospectorNode, +) +from .model_loader.zmongo_universal_model_selector_node import ( + ZMongoUniversalModelSelectorNode, +) +from .model_loader.zmongo_universal_model_loader_node import ( + ZMongoUniversalModelLoaderNode, +) +from .model_loader.zmongo_universal_adapter_node import ( + ZMongoUniversalModelAdapterNode, +) + +from .model_loader.zmongo_builtin_loader_adapter_node import ( + ZMongoBuiltInLoaderAdapterNode, +) +__all__ = [ + "register_preset_routes", + "register_zmongo_field_selector_routes", + "register_zmongo_tabular_record_view_routes", + "register_zmongo_record_editor_routes", + "ZMongoConfigNode", + "ZMongoTextFetcher", + "ZMongoOperationsNode", + "ZMongoRecordSplitter", + "ZMongoFieldSelector", + "ZRetrieverNode", + "ZMongoDatabaseBrowserNode", + "ZMongoRecordLoopNode", + "ZMongoSaveTextNode", + "ZMongoDataPassThroughNode", + "ZMongoSaveBatchTextNode", + "ZMongoSaveValueNode", + "ZMongoFlattenedFieldSelectorNode", + "ZMongoFieldSelectorNode", + "ZMongoTabularRecordViewNode", + "ZMongoRecordEditorNode", + "ZMongoModelIntrospectorNode", + "ZMongoModelCompatibilityDisplayNode", + "ZMongoUniversalModelSelectorNode", + "ZMongoUniversalModelLoaderNode", + "ZMongoUniversalModelAdapterNode", + "ZMongoBuiltInLoaderAdapterNode", + "DataProcessor", + "SafeResult", + "ZMongo", + "ZEmbedder", + "LocalVectorSearch", +] \ No newline at end of file diff --git a/gui/preset_api.py b/gui/preset_api.py new file mode 100644 index 0000000..f7215d2 --- /dev/null +++ b/gui/preset_api.py @@ -0,0 +1,121 @@ +import logging +from datetime import datetime, timezone + +from aiohttp import web + +try: + from .zmongo_manager import ZMongoManager +except Exception: + from .zmongo_manager import ZMongoManager + +try: + from .zmongo_toolbag.data_processor import DataProcessor +except Exception: + from .zmongo_toolbag.data_processor import DataProcessor + +try: + from .zmongo_toolbag.zembedder import ZEmbedder +except Exception: + from .zmongo_toolbag.zembedder import ZEmbedder + +try: + from .zmongo_toolbag.zmongo import ZMongo +except Exception: + from .zmongo_toolbag.zmongo import ZMongo + +try: + from .zmongo_toolbag.safe_result import SafeResult +except Exception: + from .zmongo_toolbag.safe_result import SafeResult +logger = logging.getLogger(__name__) + +try: + # Available in PyMongo >= 4.x + from bson import BSON, decode_file_iter + + HAVE_BSON_STREAM = True +except Exception: + HAVE_BSON_STREAM = False + + +logger = logging.getLogger(__name__) + +PRESET_COLLECTION = "comfyui_presets" +zmongo = ZMongo() + + +def _doc_to_jsonable(doc): + if not doc: + return doc + out = dict(doc) + if "_id" in out: + out["_id"] = str(out["_id"]) + return out + + +async def list_presets(request): + workflow = request.query.get("workflow") + query = {"workflow": workflow} if workflow else {} + + res = await zmongo.find_many_async(PRESET_COLLECTION, query=query, limit=1000) + if not res.success: + return web.json_response({"success": False, "error": res.error}, status=500) + + docs = [_doc_to_jsonable(d) for d in (res.data or [])] + return web.json_response({"success": True, "presets": docs}) + + +async def save_preset(request): + try: + payload = await request.json() + except Exception as e: + return web.json_response({"success": False, "error": f"Invalid JSON: {e}"}, status=400) + + name = payload.get("name") + workflow = payload.get("workflow", "default") + group = payload.get("group", "Default") + updates = payload.get("updates", {}) + + if not name: + return web.json_response({"success": False, "error": "Missing preset name"}, status=400) + + now = datetime.now(timezone.utc).isoformat() + + query = {"workflow": workflow, "name": name} + doc = { + "workflow": workflow, + "name": name, + "group": group, + "updates": updates, + "updated_at": now, + } + + res = await zmongo.insert_or_update_async(PRESET_COLLECTION, query, doc) + if not res.success: + return web.json_response({"success": False, "error": res.error}, status=500) + + return web.json_response({"success": True, "message": "Preset saved"}) + + +async def delete_preset(request): + workflow = request.match_info.get("workflow") + name = request.match_info.get("name") + + if not workflow or not name: + return web.json_response({"success": False, "error": "Missing workflow or name"}, status=400) + + res = await zmongo.delete_many_async(PRESET_COLLECTION, {"workflow": workflow, "name": name}) + if not res.success: + return web.json_response({"success": False, "error": res.error}, status=500) + + return web.json_response({"success": True, "message": "Preset deleted"}) + + +def register_preset_routes(prompt_server): + routes = prompt_server.routes + + routes.get("/zmongo/presets")(list_presets) + routes.post("/zmongo/presets")(save_preset) + routes.delete("/zmongo/presets/{workflow}/{name}")(delete_preset) + + logger.info("[ZMongo] Preset API routes registered") \ No newline at end of file diff --git a/gui/zadmin.py b/gui/zadmin.py new file mode 100644 index 0000000..0678ad6 --- /dev/null +++ b/gui/zadmin.py @@ -0,0 +1,81 @@ +import asyncio +import logging +import threading +import tkinter as tk +from tkinter import ttk + +from .zmongo_admin_tabs import ZMongoAdminTabs +from .zmongo_manager import ZMongoManager + + +class ZAdminGUI(tk.Tk): + def __init__(self): + super().__init__() + self.title("ZMongo Professional Admin System") + self.geometry("1400x900") + + self.loop = asyncio.new_event_loop() + self.loop_thread = threading.Thread( + target=self._run_loop, + name="zadmin-async-loop", + daemon=True, + ) + self.loop_thread.start() + + self.engine = self._init_engine() + self._init_ui() + + self.protocol("WM_DELETE_WINDOW", self.on_close) + + def _run_loop(self): + asyncio.set_event_loop(self.loop) + self.loop.run_forever() + + def _init_engine(self): + """ + Initialize the async singleton correctly and wait for the real instance. + """ + future = asyncio.run_coroutine_threadsafe( + ZMongoManager.get_instance(), + self.loop, + ) + return future.result() + + def _init_ui(self): + self.main_notebook = ttk.Notebook(self) + self.main_notebook.pack(expand=True, fill="both", padx=10, pady=10) + + self.tabs = ZMongoAdminTabs( + self.main_notebook, + self.engine, + self.loop, + ) + + def on_close(self): + try: + if getattr(self, "engine", None) is not None: + try: + if hasattr(self.engine, "async_client"): + self.engine.async_client.close() + if hasattr(self.engine, "sync_client"): + self.engine.sync_client.close() + except Exception as exc: + logging.error("Error closing Mongo clients: %s", exc) + + try: + type(self.engine)._instance = None + except Exception: + pass + + if hasattr(self, "loop") and self.loop.is_running(): + self.loop.call_soon_threadsafe(self.loop.stop) + + if hasattr(self, "loop_thread") and self.loop_thread.is_alive(): + self.loop_thread.join(timeout=1.0) + finally: + self.destroy() + + +if __name__ == "__main__": + app = ZAdminGUI() + app.mainloop() \ No newline at end of file diff --git a/gui/zmongo_admin_gui.py b/gui/zmongo_admin_gui.py new file mode 100644 index 0000000..b158967 --- /dev/null +++ b/gui/zmongo_admin_gui.py @@ -0,0 +1,144 @@ +import asyncio +import logging +import os +import queue +import threading +import tkinter as tk +from pathlib import Path +from tkinter import ttk + +from dotenv import load_dotenv +from motor.motor_asyncio import AsyncIOMotorClient +from pymongo import MongoClient + +from .zmongo_admin_tabs import ZMongoAdminTabs +from .zmongo_toolbag.zmongo import ZMongo + +load_dotenv(Path.home() / ".resources" / ".env") + +MONGO_URI = os.getenv("MONGO_URI", "mongodb://127.0.0.1:27017") +MONGO_DATABASE_NAME = os.getenv("MONGO_DATABASE_NAME", "test") + + +class ZMongoAdminBackend: + """ + Non-GUI backend bridge for the admin GUI. + + This is what ZMongoAdminTabs should receive. + """ + def __init__(self, loop: asyncio.AbstractEventLoop): + self.loop = loop + self.zmongo = ZMongo() + + self.async_client = AsyncIOMotorClient(MONGO_URI) + self.db = self.async_client[MONGO_DATABASE_NAME] + + self.sync_client = MongoClient(MONGO_URI) + self.sync_db = self.sync_client[MONGO_DATABASE_NAME] + + self.comfyui_to_zmongo: queue.Queue = queue.Queue() + self.zmongo_to_comfyui: queue.Queue = queue.Queue() + + def run_in_async_loop(self, async_func, *args, **kwargs): + return asyncio.run_coroutine_threadsafe(async_func(*args, **kwargs), self.loop) + + def submit_from_comfyui(self, payload: dict): + self.comfyui_to_zmongo.put(payload) + + def submit_to_comfyui(self, payload: dict): + self.zmongo_to_comfyui.put(payload) + + async def process_comfyui_payload(self, payload: dict): + """ + Store inbound ComfyUI events/results in Mongo. + """ + document = { + "source": "comfyui", + "payload_type": payload.get("type", "unknown"), + "payload": payload, + } + return await self.zmongo.insert_document("comfyui_events", document) + + def close(self): + try: + self.async_client.close() + except Exception as exc: + logging.error("Error closing async Mongo client: %s", exc) + + try: + self.sync_client.close() + except Exception as exc: + logging.error("Error closing sync Mongo client: %s", exc) + + +class ZAdminGUI(tk.Tk): + def __init__(self): + super().__init__() + self.title("ZMongo Professional Admin System") + self.geometry("1400x900") + + self.loop = asyncio.new_event_loop() + self.loop_thread = threading.Thread( + target=self._run_loop, + name="zadmin-async-loop", + daemon=True, + ) + self.loop_thread.start() + + self.engine = ZMongoAdminBackend(self.loop) + + self.main_notebook = ttk.Notebook(self) + self.main_notebook.pack(expand=True, fill="both", padx=10, pady=10) + + self.tabs = ZMongoAdminTabs(self.main_notebook, self.engine, self.loop) + + self.after(100, self._pump_bridge) + self.protocol("WM_DELETE_WINDOW", self.on_close) + + def _run_loop(self): + asyncio.set_event_loop(self.loop) + self.loop.run_forever() + + def _pump_bridge(self): + """ + Pump ComfyUI -> Mongo payloads from the Tk thread. + """ + try: + while True: + payload = self.engine.comfyui_to_zmongo.get_nowait() + future = self.engine.run_in_async_loop(self.engine.process_comfyui_payload, payload) + future.add_done_callback(self._handle_bridge_result) + except queue.Empty: + pass + + self.after(100, self._pump_bridge) + + def _handle_bridge_result(self, future): + try: + result = future.result() + self.engine.submit_to_comfyui({ + "type": "zmongo_ack", + "success": getattr(result, "success", False), + "data": getattr(result, "data", None), + "error": getattr(result, "error", None), + }) + except Exception as exc: + logging.error("Bridge task failed: %s", exc) + + def on_close(self): + try: + if getattr(self, "engine", None) is not None: + self.engine.close() + + if hasattr(self, "loop") and self.loop.is_running(): + self.loop.call_soon_threadsafe(self.loop.stop) + + if hasattr(self, "loop_thread") and self.loop_thread.is_alive(): + self.loop_thread.join(timeout=1.0) + finally: + self.destroy() + + +if __name__ == "__main__": + app = ZAdminGUI() + app.mainloop() \ No newline at end of file diff --git a/gui/zmongo_admin_tabs.py b/gui/zmongo_admin_tabs.py new file mode 100644 index 0000000..cb41f64 --- /dev/null +++ b/gui/zmongo_admin_tabs.py @@ -0,0 +1,608 @@ +import asyncio +import csv +import json +import logging +import os +import re +import subprocess +import threading +from dataclasses import dataclass +from datetime import datetime +from pathlib import Path +from tkinter import ( + BOTH, + END, + LEFT, + RIGHT, + TOP, + X, + Y, + Button, + Entry, + Frame, + Listbox, + filedialog, + ttk, +) +from tkinter.scrolledtext import ScrolledText +import tkinter as tk + +from bson import errors +from bson.objectid import ObjectId + +from .zmongo_toolbag.data_processor import DataProcessor +from .zmongo_toolbag.safe_result import SafeResult + +try: + from bson import BSON, decode_file_iter + + HAVE_BSON_STREAM = True +except Exception: + HAVE_BSON_STREAM = False + +logger = logging.getLogger(__name__) + + +@dataclass +class Pager: + limit: int = 100 + skip: int = 0 + + +class ZMongoAdminTabs: + """ + Modular admin tabs for Mongo/ZMongo management. + + Supports either: + - legacy wrapper objects exposing `.client` + - newer backend objects exposing `.zmongo` + - raw ZMongo-like objects directly + """ + + def __init__(self, parent_notebook: ttk.Notebook, zmongo_singleton, loop: asyncio.AbstractEventLoop): + self.parent = parent_notebook + self.notebook = parent_notebook + self.loop = loop + self.manager = zmongo_singleton + + self.zmongo = self._resolve_zmongo(zmongo_singleton) + self.db = self._resolve_db(zmongo_singleton, self.zmongo) + self.sync_db = getattr(zmongo_singleton, "sync_db", None) + + if self.db is None: + raise AttributeError( + "ZMongoAdminTabs requires an object with either '.client' or '.zmongo', " + "and a usable '.db' attribute." + ) + + self.db_name = self._resolve_db_name(zmongo_singleton, self.zmongo) + + self.backup_dir = Path.home() / os.getenv("MONGO_BACKUP_DIR", ".resources/backups") / self.db_name + self.backup_dir.mkdir(parents=True, exist_ok=True) + + self.cv_pager = Pager(limit=100, skip=0) + self.cv_ids_cache: list[tuple[str, object]] = [] + self.active_processes: dict[str, subprocess.Popen] = {} + + self.collection_listbox = None + self.backup_files_listbox = None + self.backup_format_combo = None + self.message_text = None + self.console = None + self.cv_collection_entry = None + self.cv_filter_entry = None + self.cv_doc_listbox = None + self.cv_json_text = None + self.cv_dotkey_entry = None + self.cv_value_entry = None + + self._build_tabs(parent_notebook) + self.run_periodic_updates() + + @staticmethod + def _resolve_zmongo(zmongo_singleton): + if hasattr(zmongo_singleton, "client"): + return zmongo_singleton.client + if hasattr(zmongo_singleton, "zmongo"): + return zmongo_singleton.zmongo + return zmongo_singleton + + @staticmethod + def _resolve_db(zmongo_singleton, zmongo_obj): + direct_db = getattr(zmongo_singleton, "db", None) + if direct_db is not None: + return direct_db + return getattr(zmongo_obj, "db", None) + + @staticmethod + def _resolve_db_name(zmongo_singleton, zmongo_obj): + db_name = getattr(zmongo_singleton, "db_name", None) + if db_name: + return db_name + db_name = getattr(zmongo_obj, "db_name", None) + if db_name: + return db_name + db = getattr(zmongo_singleton, "db", None) or getattr(zmongo_obj, "db", None) + if db is not None and hasattr(db, "name"): + return db.name + return os.getenv("MONGO_DATABASE_NAME", "test") + + def _build_tabs(self, notebook): + self.explorer_frame = ttk.Frame(notebook) + notebook.add(self.explorer_frame, text="Collection Viewer") + self._build_collection_tab(self.explorer_frame) + + self.maintenance_frame = ttk.Frame(notebook) + notebook.add(self.maintenance_frame, text="Backup & Restore") + self._build_maintenance_tab(self.maintenance_frame) + + self.system_frame = ttk.Frame(notebook) + notebook.add(self.system_frame, text="Service Runners") + self._build_system_tab(self.system_frame) + + def _build_collection_tab(self, parent: Frame): + root = ttk.Frame(parent) + root.pack(fill=BOTH, expand=True, padx=8, pady=8) + root.grid_columnconfigure(1, weight=1) + root.grid_rowconfigure(2, weight=1) + + ttk.Label(root, text="Collection:").grid(row=0, column=0, sticky="w") + self.cv_collection_entry = Entry(root, width=40) + self.cv_collection_entry.grid(row=0, column=1, sticky="w", padx=6) + + ttk.Label(root, text="Filter (JSON):").grid(row=1, column=0, sticky="w") + self.cv_filter_entry = Entry(root) + self.cv_filter_entry.insert(0, "{}") + self.cv_filter_entry.grid(row=1, column=1, sticky="we", padx=6) + + ctrl = ttk.Frame(root) + ctrl.grid(row=0, column=2, rowspan=2, sticky="ne") + Button(ctrl, text="Refresh", command=self.cv_refresh_docs_clicked).pack(side=TOP, padx=4, pady=2) + Button(ctrl, text="Load More", command=self.cv_load_more_clicked).pack(side=TOP, padx=4, pady=2) + + left = ttk.Frame(root, borderwidth=1, relief="groove") + left.grid(row=2, column=0, sticky="ns") + + self.cv_doc_listbox = Listbox(left, height=30, width=36, exportselection=False) + self.cv_doc_listbox.pack(side=LEFT, fill=Y) + self.cv_doc_listbox.bind("<>", self.cv_on_doc_select) + + scrollbar = ttk.Scrollbar(left, orient="vertical", command=self.cv_doc_listbox.yview) + scrollbar.pack(side=RIGHT, fill=Y) + self.cv_doc_listbox.config(yscrollcommand=scrollbar.set) + + right = ttk.Frame(root, borderwidth=1, relief="groove") + right.grid(row=2, column=1, columnspan=2, sticky="nswe", padx=(8, 0)) + right.grid_columnconfigure(0, weight=1) + right.grid_rowconfigure(0, weight=1) + + self.cv_json_text = ScrolledText(right, font=("Courier New", 10), wrap="none") + self.cv_json_text.grid(row=0, column=0, sticky="nswe", padx=5, pady=5) + + editor = ttk.LabelFrame(right, text="Dot-Key Update (Single Doc)") + editor.grid(row=1, column=0, sticky="we", padx=5, pady=5) + + ttk.Label(editor, text="Dot-key:").pack(side=LEFT, padx=5) + self.cv_dotkey_entry = Entry(editor, width=30) + self.cv_dotkey_entry.pack(side=LEFT, padx=5) + + ttk.Label(editor, text="Value:").pack(side=LEFT, padx=5) + self.cv_value_entry = Entry(editor, width=30) + self.cv_value_entry.pack(side=LEFT, padx=5) + + Button(editor, text="Apply $set", command=self.on_apply_dotkey_value_clicked).pack(side=RIGHT, padx=5) + + def _build_maintenance_tab(self, parent: Frame): + root = ttk.Frame(parent) + root.pack(fill=BOTH, expand=True, padx=8, pady=8) + root.grid_columnconfigure(0, weight=1) + root.grid_columnconfigure(2, weight=1) + root.grid_rowconfigure(3, weight=1) + + ttk.Label(root, text="Collections:").grid(row=0, column=0, sticky="w", padx=5) + self.collection_listbox = Listbox(root, exportselection=False, height=12) + self.collection_listbox.grid(row=1, column=0, sticky="nswe", padx=5) + self.collection_listbox.bind("<>", self.on_collection_select) + + ttk.Label(root, text="Backup Files:").grid(row=0, column=2, sticky="w", padx=5) + self.backup_files_listbox = Listbox(root, exportselection=False, height=12) + self.backup_files_listbox.grid(row=1, column=2, sticky="nswe", padx=5) + self.backup_files_listbox.bind("<>", self.on_backup_file_select) + + actions = ttk.Frame(root) + actions.grid(row=2, column=0, columnspan=4, sticky="we", pady=10) + + self.backup_format_combo = ttk.Combobox(actions, values=["JSON", "BSON", "CSV"], state="readonly", width=8) + self.backup_format_combo.current(0) + self.backup_format_combo.pack(side=LEFT, padx=5) + + Button(actions, text="Backup Selected", command=self.on_backup_selected_clicked).pack(side=LEFT, padx=5) + Button(actions, text="Restore Selected", command=self.on_restore_clicked).pack(side=LEFT, padx=5) + Button(actions, text="Browse File...", command=self.open_file_explorer).pack(side=LEFT, padx=5) + + self.message_text = ScrolledText(root, height=8, wrap="word", state="disabled") + self.message_text.grid(row=3, column=0, columnspan=4, sticky="nswe", padx=5, pady=5) + + def _build_system_tab(self, parent): + controls = ttk.LabelFrame(parent, text="Process Controls") + controls.pack(fill=X, padx=10, pady=10) + + tk.Button(controls, text="Start OCR Runner", command=lambda: self.start_service("ocr_runner.py")).pack( + side=LEFT, padx=5 + ) + tk.Button(controls, text="Start ZRetriever API", command=lambda: self.start_service("zretriever.py")).pack( + side=LEFT, padx=5 + ) + + self.console = ScrolledText(parent, height=20, bg="black", fg="white") + self.console.pack(expand=True, fill=BOTH, padx=10, pady=10) + + def start_service(self, script_path): + def _stream(): + self.notebook.after(0, lambda: self.console.insert(END, f"\n[SYSTEM] Starting {script_path}...\n")) + process = subprocess.Popen( + ["python", script_path], + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + ) + self.active_processes[script_path] = process + for line in process.stdout: + self.notebook.after(0, lambda l=line: (self.console.insert(END, l), self.console.see(END))) + + threading.Thread(target=_stream, daemon=True).start() + + def run_periodic_updates(self): + self.run_in_async_loop(self.fetch_and_update_db_info) + self.run_in_async_loop(self.fetch_and_update_collections) + self.notebook.after(30000, self.run_periodic_updates) + + async def fetch_and_update_db_info(self): + try: + await self.db.list_collection_names() + await self.db.command("dbstats") + except Exception as e: + logger.error("DB Info update failed: %s", e) + + async def fetch_and_update_collections(self): + try: + names = sorted(await self.db.list_collection_names()) + self.notebook.after(0, lambda: self._update_collection_listbox(names)) + except Exception as e: + logger.error("Collection list update failed: %s", e) + + def _update_collection_listbox(self, names): + self.collection_listbox.delete(0, END) + for name in names: + self.collection_listbox.insert(END, name) + + def run_in_async_loop(self, async_func, *args, **kwargs): + future = asyncio.run_coroutine_threadsafe(async_func(*args, **kwargs), self.loop) + future.add_done_callback(self.on_async_task_done) + return future + + def on_async_task_done(self, future): + try: + future.result() + except Exception as e: + self.log_message(f"Async Error: {e}") + + def log_message(self, message): + if self.message_text is None: + logger.error(message) + return + + def _append(): + self.message_text.config(state="normal") + self.message_text.insert(END, f"{datetime.now().strftime('%H:%M:%S')} - {message}\n") + self.message_text.config(state="disabled") + self.message_text.see(END) + + self.notebook.after(0, _append) + + def cv_refresh_docs_clicked(self): + self.cv_pager.skip = 0 + self.cv_refresh_docs(reset=True) + + def cv_load_more_clicked(self): + self.cv_pager.skip += self.cv_pager.limit + self.cv_refresh_docs(reset=False) + + def cv_refresh_docs(self, reset: bool): + collection = self.cv_collection_entry.get().strip() + if not collection: + self.log_message("Collection is required.") + return + + try: + filt = json.loads(self.cv_filter_entry.get() or "{}") + if not isinstance(filt, dict): + raise ValueError + except Exception: + filt = {} + self.log_message("Invalid JSON filter; using {}.") + + async def _fetch(): + cursor = self.db[collection].find(filt, projection={"_id": 1}).skip(self.cv_pager.skip).limit( + self.cv_pager.limit + ) + items = await cursor.to_list(length=self.cv_pager.limit) + ids = [(str(it.get("_id")), it.get("_id")) for it in items] + self.notebook.after(0, lambda: self._update_cv_listbox(ids, reset)) + + self.run_in_async_loop(_fetch) + + def _update_cv_listbox(self, ids, reset): + if reset: + self.cv_doc_listbox.delete(0, END) + self.cv_ids_cache = [] + + self.cv_ids_cache.extend(ids) + for sid, _ in ids: + self.cv_doc_listbox.insert(END, sid) + + def cv_on_doc_select(self, event=None): + selection = self.cv_doc_listbox.curselection() + if not selection: + return + + _, _id_obj = self.cv_ids_cache[selection[0]] + collection = self.cv_collection_entry.get().strip() + if not collection: + return + + async def _fetch(): + if hasattr(self.zmongo, "find_one_async"): + res = await self.zmongo.find_one_async(collection, {"_id": _id_obj}) + elif hasattr(self.zmongo, "find_one"): + maybe = self.zmongo.find_one(collection, {"_id": _id_obj}) + res = await maybe if asyncio.iscoroutine(maybe) else maybe + else: + raise AttributeError("Resolved ZMongo object does not expose find_one_async or find_one.") + + if isinstance(res, SafeResult) and res.success: + pretty = json.dumps(res.data, indent=2, default=str) + self.notebook.after( + 0, + lambda: ( + self.cv_json_text.delete("1.0", END), + self.cv_json_text.insert(END, pretty), + ), + ) + elif isinstance(res, SafeResult): + self.log_message(f"Fetch failed: {res.error}") + else: + pretty = json.dumps(res, indent=2, default=str) + self.notebook.after( + 0, + lambda: ( + self.cv_json_text.delete("1.0", END), + self.cv_json_text.insert(END, pretty), + ), + ) + + self.run_in_async_loop(_fetch) + + def on_apply_dotkey_value_clicked(self): + collection = self.cv_collection_entry.get().strip() + path = self.cv_dotkey_entry.get().strip() + raw_val = self.cv_value_entry.get().strip() + selection = self.cv_doc_listbox.curselection() + + if not all([collection, path, selection]): + self.log_message("Collection, dot-key, and selected document are required.") + return + + _id_str, _id_obj = self.cv_ids_cache[selection[0]] + + try: + value = json.loads(raw_val) + except Exception: + value = raw_val + + async def _update(): + if hasattr(self.zmongo, "update_one_async"): + res = await self.zmongo.update_one_async(collection, {"_id": _id_obj}, {"$set": {path: value}}) + elif hasattr(self.zmongo, "update_many_async"): + res = await self.zmongo.update_many_async(collection, {"_id": _id_obj}, {"$set": {path: value}}) + else: + raise AttributeError("Resolved ZMongo object does not expose update_one_async or update_many_async.") + + if isinstance(res, SafeResult) and res.success: + self.log_message(f"Updated {path} for {_id_str}") + self.cv_on_doc_select() + elif isinstance(res, SafeResult): + self.log_message(f"Update failed: {res.error}") + else: + self.log_message(f"Update completed for {_id_str}") + + self.run_in_async_loop(_update) + + def on_collection_select(self, event=None): + selection = self.collection_listbox.curselection() + if selection: + collection_name = self.collection_listbox.get(selection[0]) + self.update_backup_files_listbox(collection_name) + + def update_backup_files_listbox(self, collection_name: str): + self.backup_files_listbox.delete(0, END) + pattern = re.compile(rf"^{re.escape(collection_name)}\[\d{{14}}\]\.(json|bson|csv)$", re.I) + + try: + for file_path in self.backup_dir.iterdir(): + if file_path.is_file() and pattern.match(file_path.name): + self.backup_files_listbox.insert(END, file_path.name) + except FileNotFoundError: + self.log_message(f"Backup directory not found: {self.backup_dir}") + + def on_backup_file_select(self, event=None): + pass + + def on_backup_selected_clicked(self): + selection = self.collection_listbox.curselection() + if not selection: + self.log_message("No collection selected for backup.") + return + + collection_name = self.collection_listbox.get(selection[0]) + fmt = self.backup_format_combo.get() + self.run_in_async_loop(self.backup_collection, collection_name, fmt) + + async def backup_collection(self, collection_name: str, fmt: str): + try: + cursor = self.db[collection_name].find({}) + docs = await cursor.to_list(length=None) + timestamp = datetime.now().strftime("%Y%m%d%H%M%S") + output_path = self.backup_dir / f"{collection_name}[{timestamp}].{fmt.lower()}" + + if fmt == "JSON": + with open(output_path, "w", encoding="utf-8") as f: + json.dump(DataProcessor.to_json_compatible(docs), f, indent=2) + elif fmt == "BSON": + if not HAVE_BSON_STREAM: + raise RuntimeError("BSON streaming not available in this environment.") + with open(output_path, "wb") as f: + for doc in docs: + f.write(BSON.encode(doc)) + elif fmt == "CSV": + rows = [DataProcessor.flatten_json(DataProcessor.to_json_compatible(doc)) for doc in docs] + headers = sorted({key for row in rows for key in row.keys()}) + with open(output_path, "w", encoding="utf-8", newline="") as f: + writer = csv.DictWriter(f, fieldnames=headers) + writer.writeheader() + for row in rows: + writer.writerow(row) + else: + raise ValueError(f"Unsupported backup format: {fmt}") + + self.log_message(f"Backed up {collection_name} to {output_path.name}") + self.notebook.after(0, lambda: self.update_backup_files_listbox(collection_name)) + except Exception as e: + self.log_message(f"Backup Error: {e}") + + def on_restore_clicked(self): + selection = self.collection_listbox.curselection() + file_selection = self.backup_files_listbox.curselection() + + if not selection: + self.log_message("No collection selected for restore.") + return + if not file_selection: + self.log_message("No backup file selected for restore.") + return + + collection_name = self.collection_listbox.get(selection[0]) + filename = self.backup_files_listbox.get(file_selection[0]) + + self.run_in_async_loop(self.restore_collection, collection_name, filename) + + async def restore_collection(self, collection_name: str, filename: str): + restore_path = self.backup_dir / filename + if not restore_path.exists(): + self.log_message(f"Restore file not found: {restore_path}") + return + + try: + ext = restore_path.suffix.lower() + docs = [] + + if ext == ".json": + with open(restore_path, "r", encoding="utf-8") as f: + data = json.load(f) + docs = data if isinstance(data, list) else [data] + + elif ext == ".bson": + if not HAVE_BSON_STREAM: + raise RuntimeError("BSON restore not supported in this environment.") + with open(restore_path, "rb") as f: + docs = list(decode_file_iter(f)) + + elif ext == ".csv": + with open(restore_path, "r", encoding="utf-8", newline="") as f: + reader = csv.DictReader(f) + docs = [self._unflatten_csv_row(row) for row in reader] + else: + raise ValueError(f"Unsupported restore file type: {ext}") + + normalized_docs = [self._normalize_doc_ids(doc) for doc in docs] + if not normalized_docs: + self.log_message("Restore file contained no documents.") + return + + await self.db[collection_name].delete_many({}) + await self.db[collection_name].insert_many(normalized_docs) + self.log_message(f"Restored {len(normalized_docs)} documents into {collection_name}") + except Exception as e: + self.log_message(f"Restore Error: {e}") + + def open_file_explorer(self): + filepath = filedialog.askopenfilename( + initialdir=str(self.backup_dir), + title="Select a Backup File", + filetypes=[ + ("All supported", "*.json *.bson *.csv"), + ("JSON files", "*.json"), + ("BSON files", "*.bson"), + ("CSV files", "*.csv"), + ], + ) + if not filepath: + return + + file_path_obj = Path(filepath) + filename = file_path_obj.name + + self.backup_files_listbox.selection_clear(0, END) + existing = list(self.backup_files_listbox.get(0, END)) + if filename not in existing: + self.backup_files_listbox.insert(END, filename) + + index = list(self.backup_files_listbox.get(0, END)).index(filename) + self.backup_files_listbox.selection_set(index) + self.backup_files_listbox.see(index) + + @staticmethod + def _normalize_doc_ids(doc: dict): + if isinstance(doc, dict) and "_id" in doc and isinstance(doc["_id"], str) and ObjectId.is_valid(doc["_id"]): + try: + doc["_id"] = ObjectId(doc["_id"]) + except errors.InvalidId: + pass + return doc + + @staticmethod + def _unflatten_csv_row(row: dict) -> dict: + root = {} + + def set_path(container, parts, value): + if not parts: + return value + + key = parts[0] + is_index = key.isdigit() + + if is_index: + index = int(key) + if not isinstance(container, list): + container = [] + while len(container) <= index: + container.append({}) + container[index] = set_path(container[index], parts[1:], value) + return container + + if not isinstance(container, dict): + container = {} + + container[key] = set_path(container.get(key, {}), parts[1:], value) + return container + + for k, v in row.items(): + if v in (None, ""): + continue + try: + parsed = json.loads(v) + except Exception: + parsed = v + root = set_path(root, k.split("."), parsed) + + return root \ No newline at end of file diff --git a/gui/zmongo_manager.py b/gui/zmongo_manager.py new file mode 100644 index 0000000..9fa7d9c --- /dev/null +++ b/gui/zmongo_manager.py @@ -0,0 +1,468 @@ +# zmongo_manager.py +import asyncio +import csv +import json +import logging +import os +import threading +from dataclasses import dataclass +from datetime import datetime +from pathlib import Path +from tkinter import ttk, Tk, Listbox, Entry, Button, END, BOTH, LEFT, X +from tkinter.scrolledtext import ScrolledText + +from bson import json_util +from dotenv import load_dotenv +from motor.motor_asyncio import AsyncIOMotorClient +from pymongo import MongoClient + +# Handle dynamic imports for the toolbag +try: + from .zmongo_toolbag.data_processor import DataProcessor + from .zmongo_toolbag.zmongo import ZMongo + from .zmongo_toolbag.safe_result import SafeResult +except (ImportError, ValueError): + from .zmongo_toolbag.data_processor import DataProcessor + from .zmongo_toolbag.zmongo import ZMongo + from .zmongo_toolbag.safe_result import SafeResult + +# BSON Streaming support check +try: + from bson import BSON, decode_file_iter + + HAVE_BSON_STREAM = True +except ImportError: + HAVE_BSON_STREAM = False + +# --- Configuration --- +load_dotenv(Path.home() / ".resources" / ".env") +logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') + +MONGO_URI = os.getenv("MONGO_URI", "mongodb://127.0.0.1:27017") +MONGO_DATABASE_NAME = os.getenv("MONGO_DATABASE_NAME", "test") +MONGO_BACKUP_DIR_REL = os.getenv("MONGO_BACKUP_DIR", '.resources/mongo_backups') + + +@dataclass +class Pager: + limit: int = 100 + skip: int = 0 + + +class ZMongoManager(Tk): + _instance = None + _instance_lock = asyncio.Lock() + + def __init__(self, loop: asyncio.AbstractEventLoop = None): + super().__init__() + self.zmongo = ZMongo() + self.title("ZMongo System Manager") + self.geometry("1400x900") + + # Fallback to current loop if not provided + self.loop = loop or asyncio.get_event_loop() + self.db_name = MONGO_DATABASE_NAME + self.backup_dir = Path.home() / MONGO_BACKUP_DIR_REL / MONGO_DATABASE_NAME + self.backup_dir.mkdir(parents=True, exist_ok=True) + + try: + self.async_client = AsyncIOMotorClient(MONGO_URI) + self.db = self.async_client[self.db_name] + self.sync_client = MongoClient(MONGO_URI) + self.sync_db = self.sync_client[self.db_name] + except Exception as e: + logging.error(f"Failed to connect to MongoDB: {e}") + self.destroy() + return + + self.cv_pager = Pager() + self.cv_ids_cache = [] + + self._create_widgets() + self.run_periodic_updates() + + @classmethod + async def get_instance(cls): + """ + Async singleton accessor. Requires no arguments. + Detects the running loop automatically. + """ + if cls._instance is not None: + return cls._instance + + async with cls._instance_lock: + if cls._instance is None: + # Get the loop currently running this coroutine + loop = asyncio.get_running_loop() + cls._instance = cls(loop) + return cls._instance + + def cv_refresh_docs_clicked(self): + coll = self.cv_collection_entry.get().strip() + if not coll: + self.log_message("Please select a collection first.") + return + + raw_filter = self.cv_filter_entry.get().strip() or "{}" + try: + filt = json.loads(raw_filter) + if not isinstance(filt, dict): + self.log_message("Filter must be a JSON object.") + return + except json.JSONDecodeError as e: + self.log_message(f"Invalid JSON filter: {e}") + return + + async def _fetch(): + cursor = self.db[coll].find(filt).limit(100) + docs = await cursor.to_list(length=100) + + rows = [] + for doc in docs: + oid = doc.get("_id") + oid_str = str(oid) if oid is not None else "" + + preview = ( + doc.get("name") + or doc.get("title") + or doc.get("username") + or doc.get("email") + or doc.get("metadata.title") + or "" + ) + + if not preview: + flat = DataProcessor.flatten_json(DataProcessor.normalize_objectid(doc)) + preview_parts = [] + for key in ("name", "title", "username", "email", "metadata.title", "metadata.source_url"): + val = flat.get(key) + if val: + preview_parts.append(f"{key}={val}") + preview = " | ".join(preview_parts[:2]) + + label = oid_str if not preview else f"{oid_str} | {preview}" + rows.append((label, oid)) + + return rows + + def _on_done(f): + try: + pairs = f.result() + except Exception as e: + self.after(0, lambda: self.log_message(f"Refresh failed for '{coll}': {e}")) + return + + def _update(): + self.cv_doc_listbox.delete(0, END) + self.cv_ids_cache = pairs + self.cv_json_text.delete("1.0", END) + + for label, _ in pairs: + self.cv_doc_listbox.insert(END, label) + + self.log_message(f"Loaded {len(pairs)} record(s) from '{coll}'.") + + self.after(0, _update) + + self.run_in_async_loop(_fetch).add_done_callback(_on_done) + + def cv_on_doc_select(self, event=None): + idx = self.cv_doc_listbox.curselection() + if not idx: + return + + coll = self.cv_collection_entry.get().strip() + if not coll: + self.log_message("No collection selected.") + return + + if idx[0] >= len(self.cv_ids_cache): + self.log_message("Selected record index is out of range.") + return + + _, oid = self.cv_ids_cache[idx[0]] + + async def _get(): + return await self.db[coll].find_one({"_id": oid}) + + def _on_done(f): + try: + doc = f.result() + except Exception as e: + self.after(0, lambda: self.log_message(f"Document load failed for '{coll}': {e}")) + return + + def _show(): + self.cv_json_text.delete("1.0", END) + if doc is None: + self.cv_json_text.insert(END, "{}") + self.log_message("Document not found.") + return + self.cv_json_text.insert(END, json_util.dumps(doc, indent=2)) + + self.after(0, _show) + + self.run_in_async_loop(_get).add_done_callback(_on_done) + + async def fetch_and_update_collections(self): + names = sorted(await self.db.list_collection_names()) + + def _update(): + self.collection_listbox.delete(0, END) + if hasattr(self, "cv_collection_listbox"): + self.cv_collection_listbox.delete(0, END) + + for n in names: + self.collection_listbox.insert(END, n) + if hasattr(self, "cv_collection_listbox"): + self.cv_collection_listbox.insert(END, n) + + self.after(0, _update) + + # ---------- UI BUILD ---------- + def _create_widgets(self): + self.notebook = ttk.Notebook(self) + self.notebook.pack(expand=True, fill=BOTH, padx=10, pady=10) + + tabs = { + "info": ttk.Frame(self.notebook), + "maint": ttk.Frame(self.notebook), + "viewer": ttk.Frame(self.notebook), + "system": ttk.Frame(self.notebook) + } + + self.notebook.add(tabs["info"], text='Database Info') + self.notebook.add(tabs["maint"], text='Backup & Restore') + self.notebook.add(tabs["viewer"], text='Collection Viewer / Editor') + self.notebook.add(tabs["system"], text='System Runner') + + self.db_info_text = ScrolledText(tabs["info"], wrap="word", font=("Courier New", 10)) + self.db_info_text.pack(expand=True, fill=BOTH, padx=5, pady=5) + + self._build_maintenance_tab(tabs["maint"]) + self._build_collection_tab(tabs["viewer"]) + + def _build_maintenance_tab(self, parent): + root = ttk.Frame(parent) + root.pack(fill=BOTH, expand=True, padx=8, pady=8) + root.grid_columnconfigure((1, 3), weight=1) + root.grid_rowconfigure(5, weight=1) + + # Collection List + self.collection_listbox = Listbox(root, exportselection=False, height=12) + self.collection_listbox.grid(row=1, column=0, sticky="nswe", padx=5) + self.collection_listbox.bind('<>', self.on_collection_select) + + # Backup Files List + self.backup_files_listbox = Listbox(root, exportselection=False, height=12) + self.backup_files_listbox.grid(row=1, column=2, sticky="nswe", padx=5) + self.backup_files_listbox.bind('<>', self.on_backup_file_select) + + # Entries + self.selected_collection_entry = Entry(root, state='readonly') + self.selected_collection_entry.grid(row=1, column=1, sticky="we", padx=5) + self.selected_backup_entry = Entry(root, state='readonly') + self.selected_backup_entry.grid(row=1, column=3, sticky="we", padx=5) + + # Actions + actions = ttk.Frame(root) + actions.grid(row=2, column=1, columnspan=3, sticky="we", pady=10) + self.backup_format_combo = ttk.Combobox(actions, state='readonly', values=["JSON", "BSON", "CSV"], width=8) + self.backup_format_combo.current(0) + self.backup_format_combo.pack(side=LEFT, padx=5) + + Button(actions, text='Backup Selected', command=self.on_backup_selected_clicked).pack(side=LEFT, padx=5) + Button(actions, text='Restore Selected', command=self.on_restore_clicked).pack(side=LEFT, padx=5) + Button(actions, text='Browse...', command=self.open_file_explorer).pack(side=LEFT, padx=5) + + self.message_text = ScrolledText(root, height=8, state='disabled') + self.message_text.grid(row=5, column=0, columnspan=4, sticky="nswe", padx=5) + + def _build_collection_tab(self, parent): + root = ttk.Frame(parent) + root.pack(fill=BOTH, expand=True, padx=8, pady=8) + root.grid_columnconfigure(1, weight=1) + root.grid_rowconfigure(2, weight=1) + + # --- Left side: collection list + document list --- + left = ttk.Frame(root) + left.grid(row=0, column=0, rowspan=3, sticky="ns", padx=5, pady=5) + left.grid_rowconfigure(1, weight=1) + left.grid_rowconfigure(3, weight=1) + + ttk.Label(left, text="Collections").grid(row=0, column=0, sticky="w") + self.cv_collection_listbox = Listbox(left, exportselection=False, height=12, width=35) + self.cv_collection_listbox.grid(row=1, column=0, sticky="nswe", pady=(0, 8)) + self.cv_collection_listbox.bind("<>", self.cv_on_collection_select) + + ttk.Label(left, text="Records").grid(row=2, column=0, sticky="w") + self.cv_doc_listbox = Listbox(left, exportselection=False, width=45) + self.cv_doc_listbox.grid(row=3, column=0, sticky="nswe") + self.cv_doc_listbox.bind("<>", self.cv_on_doc_select) + + # --- Top controls --- + ttk.Label(root, text="Selected Collection").grid(row=0, column=1, sticky="w", padx=5) + self.cv_collection_entry = Entry(root) + self.cv_collection_entry.grid(row=0, column=1, sticky="we", padx=5) + + ttk.Label(root, text="Filter JSON").grid(row=1, column=1, sticky="w", padx=5) + self.cv_filter_entry = Entry(root) + self.cv_filter_entry.insert(0, "{}") + self.cv_filter_entry.grid(row=1, column=1, sticky="we", padx=5) + + btn_frame = ttk.Frame(root) + btn_frame.grid(row=0, column=2, rowspan=2, sticky="ns", padx=5) + Button(btn_frame, text="Refresh Records", command=self.cv_refresh_docs_clicked).pack(fill=X, pady=2) + Button(btn_frame, text="Apply $set", command=self.on_apply_dotkey_value_clicked).pack(fill=X, pady=2) + + # --- JSON viewer --- + self.cv_json_text = ScrolledText(root, font=("Courier New", 10)) + self.cv_json_text.grid(row=2, column=1, columnspan=2, sticky="nswe", padx=5, pady=5) + + # ---------- Logic ---------- + def run_in_async_loop(self, async_func, *args, **kwargs): + return asyncio.run_coroutine_threadsafe(async_func(*args, **kwargs), self.loop) + + def _handle_future_result(self, future): + try: + future.result() + except Exception as e: + self.log_message(f"Async Error: {e}") + + def log_message(self, message): + def _append(): + self.message_text.config(state='normal') + self.message_text.insert("end", f"{datetime.now().strftime('%H:%M:%S')} - {message}\n") + self.message_text.config(state='disabled') + self.message_text.see("end") + + self.after(0, _append) + + async def backup_collection(self, collection_name, fmt): + try: + docs = await self.db[collection_name].find({}).to_list(length=None) + ts = datetime.now().strftime("%Y%m%d%H%M%S") + path = self.backup_dir / f"{collection_name}[{ts}].{fmt.lower()}" + + if fmt == "JSON": + path.write_text(json_util.dumps(docs, indent=2)) + elif fmt == "BSON" and HAVE_BSON_STREAM: + with open(path, 'wb') as f: + for d in docs: f.write(BSON.encode(d)) + elif fmt == "CSV": + flat = [DataProcessor.flatten_json(d) for d in docs] + if flat: + keys = sorted(list(set().union(*(d.keys() for d in flat)))) + with open(path, 'w', newline='') as f: + writer = csv.DictWriter(f, fieldnames=keys) + writer.writeheader() + writer.writerows(flat) + + self.log_message(f"Backup saved: {path.name}") + self.after(0, lambda: self.update_backup_files_listbox(collection_name)) + except Exception as e: + self.log_message(f"Backup failed: {e}") + + # ---------- Event Handlers ---------- + def cv_on_collection_select(self, event=None): + idx = self.cv_collection_listbox.curselection() + if not idx: + return + + name = self.cv_collection_listbox.get(idx[0]) + + self.cv_collection_entry.delete(0, END) + self.cv_collection_entry.insert(0, name) + + # Keep backup tab collection selection synchronized + self.selected_collection_entry.config(state='normal') + self.selected_collection_entry.delete(0, END) + self.selected_collection_entry.insert(0, name) + self.selected_collection_entry.config(state='readonly') + + self._select_listbox_value(self.collection_listbox, name) + self.update_backup_files_listbox(name) + self.cv_refresh_docs_clicked() + + def on_collection_select(self, event=None): + idx = self.collection_listbox.curselection() + if not idx: + return + + name = self.collection_listbox.get(idx[0]) + + self.selected_collection_entry.config(state='normal') + self.selected_collection_entry.delete(0, END) + self.selected_collection_entry.insert(0, name) + self.selected_collection_entry.config(state='readonly') + + # Keep collection viewer synchronized + self.cv_collection_entry.delete(0, END) + self.cv_collection_entry.insert(0, name) + + if hasattr(self, "cv_collection_listbox"): + self._select_listbox_value(self.cv_collection_listbox, name) + + self.update_backup_files_listbox(name) + self.cv_refresh_docs_clicked() + + def on_backup_selected_clicked(self): + name = self.selected_collection_entry.get() + fmt = self.backup_format_combo.get() + if name: self.run_in_async_loop(self.backup_collection, name, fmt) + + def _select_listbox_value(self, listbox, value): + try: + items = listbox.get(0, END) + if value in items: + idx = items.index(value) + listbox.selection_clear(0, END) + listbox.selection_set(idx) + listbox.activate(idx) + listbox.see(idx) + except Exception as e: + self.log_message(f"List selection sync failed: {e}") + + # --- Cleanup and Periodic --- + def update_backup_files_listbox(self, name): + self.backup_files_listbox.delete(0, END) + if self.backup_dir.exists(): + for f in self.backup_dir.glob(f"{name}[*].*"): + self.backup_files_listbox.insert(END, f.name) + + def run_periodic_updates(self): + self.run_in_async_loop(self.fetch_and_update_collections) + self.after(60000, self.run_periodic_updates) + + def on_closing(self): + self.async_client.close() + self.sync_client.close() + type(self)._instance = None + self.loop.call_soon_threadsafe(self.loop.stop) + self.destroy() + + # Stub methods for missing logic from snippet + def on_restore_clicked(self): + pass + + def open_file_explorer(self): + pass + + def on_backup_file_select(self, event): + pass + + def on_apply_dotkey_value_clicked(self): + pass + + +def main(): + loop = asyncio.new_event_loop() + threading.Thread(target=loop.run_forever, daemon=True).start() + + app = ZMongoManager(loop) + app.protocol("WM_DELETE_WINDOW", app.on_closing) + app.mainloop() + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/gui/zmongo_wrapper.py b/gui/zmongo_wrapper.py new file mode 100644 index 0000000..9c8c673 --- /dev/null +++ b/gui/zmongo_wrapper.py @@ -0,0 +1,36 @@ +import json + +from .zmongo_manager import ZMongoManager +from .zmongo_toolbag.data_processor import DataProcessor + + +class ZMongoOperationsNode: + """Handles unique data ingestion and maintenance tasks.""" + + @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 execute_op(self, collection, data_json, operation_type): + manager = ZMongoManager.get_instance() + data = json.loads(data_json) + records = data if isinstance(data, list) else [data] + + if operation_type == "Add Unique (No Update)": + # Implements 'check-before-insert' via async bulk_write + res = manager.client.run_sync(manager.client.bulk_write_async, collection, records) + return (DataProcessor.to_json(res.data), res.data.get("inserted_count", 0)) + + res = manager.client.insert_many(collection, records) + return (DataProcessor.to_json(res.data), len(res.data.get("inserted_ids", []))) \ No newline at end of file diff --git a/nodes/__init__.py b/nodes/__init__.py new file mode 100644 index 0000000..222b8bf --- /dev/null +++ b/nodes/__init__.py @@ -0,0 +1,28 @@ +# 1. SPECIALIZED MODEL NODE IMPORTS 🧠 +# Using absolute imports ensures these are found during standalone tests. +from ..nodes.model_loader.zmongo_model_compatibility_display_node import ZMongoModelCompatibilityDisplayNode +from ..nodes.model_loader.zmongo_model_introspector_node import ZMongoModelIntrospectorNode +from ..nodes.model_loader.zmongo_universal_model_selector_node import ZMongoUniversalModelSelectorNode +from ..nodes.model_loader.zmongo_universal_model_loader_node import ZMongoUniversalModelLoaderNode +from ..nodes.model_loader.zmongo_universal_adapter_node import ZMongoUniversalModelAdapterNode +from ..nodes.model_loader.zmongo_builtin_loader_adapter_node import ZMongoBuiltInLoaderAdapterNode + +# 2. SUB-PACKAGE MAPPINGS πŸ“‹ +# Defining these here allows the parent 'nodes' folder to simply import and merge them. +NODE_CLASS_MAPPINGS = { + "ZMongoUniversalModelSelectorNode": ZMongoUniversalModelSelectorNode, + "ZMongoModelIntrospectorNode": ZMongoModelIntrospectorNode, + "ZMongoModelCompatibilityDisplayNode": ZMongoModelCompatibilityDisplayNode, + "ZMongoUniversalModelLoaderNode": ZMongoUniversalModelLoaderNode, + "ZMongoUniversalModelAdapterNode": ZMongoUniversalModelAdapterNode, + "ZMongoBuiltInLoaderAdapterNode": ZMongoBuiltInLoaderAdapterNode, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "ZMongoUniversalModelSelectorNode": "🧩 ZMongo Universal Model Selector", + "ZMongoModelIntrospectorNode": "🧠 ZMongo Model Introspector", + "ZMongoModelCompatibilityDisplayNode": "βœ… ZMongo Model Compatibility Display", + "ZMongoUniversalModelLoaderNode": "πŸ“¦ ZMongo Universal Model Loader", + "ZMongoUniversalModelAdapterNode": "πŸ”Œ ZMongo Universal Model Adapter", + "ZMongoBuiltInLoaderAdapterNode": "βš™οΈ ZMongo Built-In Loader Adapter", +} \ No newline at end of file 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/requirements.txt b/requirements.txt index 3f1b0ec..b456387 100644 --- a/requirements.txt +++ b/requirements.txt @@ -17,4 +17,5 @@ FlagEmbedding transformers pymilvus[model] FlagEmbedding -pytest \ No newline at end of file +pytest +quart \ No newline at end of file diff --git a/server/__init__.py b/server/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/server/server.py b/server/server.py new file mode 100644 index 0000000..4db7939 --- /dev/null +++ b/server/server.py @@ -0,0 +1,28 @@ +import logging +from aiohttp import web + +logger = logging.getLogger(__name__) + +class PromptServer: + """A mock of the ComfyUI PromptServer for standalone testing.""" + instance = None + + def __init__(self): + self.app = web.Application() + self.routes = web.RouteTableDef() + # Initialize the singleton instance + PromptServer.instance = self + + def add_routes(self): + """Finalize and add the collected routes to the app.""" + self.app.add_routes(self.routes) + + async def start(self, host="127.0.0.1", port=8188): + """Starts a real local web server.""" + self.add_routes() + runner = web.AppRunner(self.app) + await runner.setup() + site = web.TCPSite(runner, host, port) + print(f"πŸš€ ZMongo Demo Server started at http://{host}:{port}") + print(f"πŸ”— Test Link: http://{host}:{port}/zmongo/flattened_fields?collection_name=test") + await site.start() \ No newline at end of file diff --git a/standalone_demo.py b/standalone_demo.py new file mode 100644 index 0000000..4a7f407 --- /dev/null +++ b/standalone_demo.py @@ -0,0 +1,50 @@ +import os +import sys +import logging +import asyncio +from aiohttp import web + +# 1. PATH CONFIGURATION πŸ› οΈ +# Ensures we can find 'nodes' and 'zmongo_toolbag' +ROOT_PATH = os.path.dirname(os.path.abspath(__file__)) +if ROOT_PATH not in sys.path: + sys.path.append(ROOT_PATH) + +# 2. WEB ASSET REGISTRATION 🌐 +WEB_DIRECTORY = "./web" + +logging.basicConfig(level=logging.INFO) +logger = logging.getLogger("ZMongo-Demo") + +# 3. IMPORTS +# We import the routes and the server class we just made +from server import PromptServer +from nodes.zmongo_field_selector_api import register_zmongo_field_selector_routes +from nodes.zmongo_record_editor_api import register_zmongo_record_editor_routes +from nodes.zmongo_tabular_record_view_api import register_zmongo_tabular_record_view_routes + + +# 4. SERVER API REGISTRATION 🚦 +def setup_server(): + server = PromptServer() + + # Register your custom routes with the server instance + register_zmongo_field_selector_routes(server) + register_zmongo_record_editor_routes(server) + register_zmongo_tabular_record_view_routes(server) + + return server + + +# 5. RUNNABLE DEMO πŸš€ +if __name__ == "__main__": + # Create the server environment + demo_server = setup_server() + + # Run the async loop + loop = asyncio.get_event_loop() + try: + loop.run_until_complete(demo_server.start()) + loop.run_forever() + except KeyboardInterrupt: + print("\nStopping ZMongo Demo Server...") \ No newline at end of file diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/test_data_processor_to_json_compatible.py b/tests/test_data_processor_to_json_compatible.py index 05dd903..3044119 100644 --- a/tests/test_data_processor_to_json_compatible.py +++ b/tests/test_data_processor_to_json_compatible.py @@ -1,13 +1,29 @@ +import os +import sys + +# 1. LIVE ENVIRONMENT MIMICRY 🎭 +# This finds the 'ComfyUI-ZMongo' folder and adds it to the search path. +# It mimics exactly what your root __init__.py does when ComfyUI starts. +CURRENT_DIR = os.path.dirname(os.path.abspath(__file__)) +PROJECT_ROOT = os.path.abspath(os.path.join(CURRENT_DIR, "..")) + +if PROJECT_ROOT not in sys.path: + sys.path.append(PROJECT_ROOT) + +# 2. ABSOLUTE PACKAGE IMPORT πŸ“¦ +# Now we can import DataProcessor as if it were a global package. +from zmongo_toolbag.data_processor import DataProcessor + +# --- Rest of your original imports --- import datetime import math import uuid from collections import deque from decimal import Decimal - import pytest from bson import ObjectId -from zmongo_toolbag.data_processor import DataProcessor +# ... rest of your test classes and functions ... class DummyModelDump: diff --git a/zmongo_toolbag/__init__.py b/zmongo_toolbag/__init__.py new file mode 100644 index 0000000..67e98f7 --- /dev/null +++ b/zmongo_toolbag/__init__.py @@ -0,0 +1,21 @@ +from .local_vector_search import LocalVectorSearch +from .zembedder import ZEmbedder +from .data_processor import DataProcessor +from .safe_result import SafeResult +from .zmongo import ZMongo + +__all__ = [ + DataProcessor, + SafeResult, + ZMongo, + ZEmbedder, + LocalVectorSearch, +] + +# __all__ = [ +# "DataProcessor", +# "SafeResult", +# "ZMongo", +# "ZEmbedder", +# "LocalVectorSearch", +# ] diff --git a/zmongo_toolbag/data_processor.py b/zmongo_toolbag/data_processor.py new file mode 100644 index 0000000..1261f55 --- /dev/null +++ b/zmongo_toolbag/data_processor.py @@ -0,0 +1,434 @@ +import base64 +import datetime +import html +import json +import logging +import math +import re +import uuid +from collections import deque +from decimal import Decimal +from typing import Any, Dict, List, Optional, Sequence, Tuple, Union + +try: + from bson import ObjectId +except ImportError: # pragma: no cover - bson should normally be available + from bson.objectid import ObjectId # type: ignore + +logger = logging.getLogger(__name__) + + +class DataProcessor: + """Centralized data processing, normalization, flattening, and text helpers.""" + + # ------------------------------------------------------------------ + # Entity & index helpers (legacy compatibility) + # ------------------------------------------------------------------ + @staticmethod + def get_entity_name(data_row: Sequence[Any]) -> Tuple[str, str]: + """Build a fallback entity name from the populated values in a row.""" + last_index = DataProcessor.get_index_last_non_nan(data_row) + entity_name = DataProcessor.get_string(data_row, 0, last_index) + return "Entity", entity_name + + @staticmethod + def get_index_last_non_excluded( + name_row: Sequence[Any], + excluded_parts: Optional[set] = None, + ) -> Optional[int]: + """Return the last non-empty index whose value is not excluded.""" + excluded_parts = excluded_parts or set() + last_value_index = DataProcessor.get_index_last_non_nan(name_row) + + for i in range(last_value_index, -1, -1): + try: + value = name_row[i] + if value not in excluded_parts: + return i + except Exception as exc: # pragma: no cover - defensive logging + logger.error("Error in get_index_last_non_excluded: %s", exc) + return None + + @staticmethod + def get_index_last_non_nan(data: Sequence[Any]) -> int: + """Return the last index whose value is not None and not NaN-like.""" + for i in range(len(data) - 1, -1, -1): + value = data[i] + if value is not None and str(value).lower() != "nan": + return i + return 0 + + @staticmethod + def get_string(data: Sequence[Any], start: int, end: int) -> str: + """Join a slice of values into a single space-delimited string.""" + if not data: + return "" + end = min(end, len(data) - 1) + start = max(start, 0) + if end < start: + return "" + parts = data[start : end + 1] + return " ".join(str(part) for part in parts if part is not None) + + # ------------------------------------------------------------------ + # BSON / datetime normalization + # ------------------------------------------------------------------ + @staticmethod + def normalize_objectid(obj: Any) -> Any: + """Legacy alias retained for backward compatibility.""" + return DataProcessor.to_json_compatible(obj) + + @staticmethod + def to_json_compatible( + data: Any, + _seen: Optional[set] = None, + *, + max_depth: int = 50, + _depth: int = 0, + ) -> Any: + """Recursively convert arbitrary Python objects into JSON-compatible data.""" + if _seen is None: + _seen = set() + + if _depth > max_depth: + return {"__truncated__": f"max_depth_exceeded:{max_depth}"} + + if data is None or isinstance(data, (bool, int, str)): + return data + + if isinstance(data, float): + if math.isnan(data): + return "NaN" + if math.isinf(data): + return "Infinity" if data > 0 else "-Infinity" + return data + + if isinstance(data, ObjectId): + return str(data) + + if isinstance(data, uuid.UUID): + return str(data) + + if isinstance(data, Decimal): + try: + if data == data.to_integral_value(): + return int(data) + return float(data) + except Exception: + return str(data) + + if isinstance(data, (datetime.datetime, datetime.date, datetime.time)): + try: + return data.isoformat() + except Exception: + return str(data) + + if isinstance(data, re.Pattern): + return data.pattern + + if isinstance(data, (bytes, bytearray, memoryview)): + raw = bytes(data) + try: + return raw.decode("utf-8") + except Exception: + return { + "__type__": "bytes", + "encoding": "base64", + "data": base64.b64encode(raw).decode("ascii"), + } + + if isinstance(data, BaseException): + return { + "__type__": data.__class__.__name__, + "message": str(data), + "args": DataProcessor.to_json_compatible( + list(getattr(data, "args", [])), + _seen=_seen.copy(), + max_depth=max_depth, + _depth=_depth + 1, + ), + } + + needs_cycle_tracking = isinstance( + data, + (dict, list, tuple, set, frozenset, deque), + ) or hasattr(data, "__dict__") or hasattr(data, "model_dump") or hasattr(data, "dict") + + if needs_cycle_tracking: + obj_id = id(data) + if obj_id in _seen: + return {"__circular_reference__": type(data).__name__} + _seen.add(obj_id) + + if isinstance(data, dict): + converted: Dict[str, Any] = {} + for key, value in data.items(): + if isinstance(key, (str, int, float, bool)) or key is None: + safe_key = str(key) + else: + safe_key = str( + DataProcessor.to_json_compatible( + key, + _seen=_seen.copy(), + max_depth=max_depth, + _depth=_depth + 1, + ) + ) + + converted[safe_key] = DataProcessor.to_json_compatible( + value, + _seen=_seen.copy(), + max_depth=max_depth, + _depth=_depth + 1, + ) + return converted + + if isinstance(data, (list, tuple, set, frozenset, deque)): + iterable = list(data) + if isinstance(data, (set, frozenset)): + try: + iterable = sorted(iterable, key=lambda x: repr(x)) + except Exception: + pass + + return [ + DataProcessor.to_json_compatible( + item, + _seen=_seen.copy(), + max_depth=max_depth, + _depth=_depth + 1, + ) + for item in iterable + ] + + try: + import numpy as np # type: ignore + + if isinstance(data, np.generic): + return DataProcessor.to_json_compatible( + data.item(), + _seen=_seen.copy(), + max_depth=max_depth, + _depth=_depth + 1, + ) + + if isinstance(data, np.ndarray): + return DataProcessor.to_json_compatible( + data.tolist(), + _seen=_seen.copy(), + max_depth=max_depth, + _depth=_depth + 1, + ) + except Exception: + pass + + try: + import pandas as pd # type: ignore + + if isinstance(data, pd.DataFrame): + return DataProcessor.to_json_compatible( + data.to_dict(orient="records"), + _seen=_seen.copy(), + max_depth=max_depth, + _depth=_depth + 1, + ) + + if isinstance(data, pd.Series): + return DataProcessor.to_json_compatible( + data.to_dict(), + _seen=_seen.copy(), + max_depth=max_depth, + _depth=_depth + 1, + ) + except Exception: + pass + + if hasattr(data, "model_dump") and callable(getattr(data, "model_dump")): + try: + dumped = data.model_dump() + return DataProcessor.to_json_compatible( + dumped, + _seen=_seen.copy(), + max_depth=max_depth, + _depth=_depth + 1, + ) + except Exception: + pass + + if hasattr(data, "dict") and callable(getattr(data, "dict")): + try: + dumped = data.dict() + return DataProcessor.to_json_compatible( + dumped, + _seen=_seen.copy(), + max_depth=max_depth, + _depth=_depth + 1, + ) + except Exception: + pass + + if hasattr(data, "__dict__"): + try: + attrs = { + key: value + for key, value in vars(data).items() + if not key.startswith("_") + } + return DataProcessor.to_json_compatible( + attrs, + _seen=_seen.copy(), + max_depth=max_depth, + _depth=_depth + 1, + ) + except Exception: + pass + + try: + return str(data) + except Exception: + return f"" + + # ------------------------------------------------------------------ + # JSON serialization + # ------------------------------------------------------------------ + @staticmethod + def to_json(data: Any, indent: Optional[int] = None) -> str: + """Serialize arbitrary data to JSON, coercing unsupported values.""" + try: + normalized = DataProcessor.to_json_compatible(data) + return json.dumps(normalized, indent=indent, ensure_ascii=False) + except Exception as exc: + logger.error("Error serializing to JSON: %s", exc) + return json.dumps({"error": str(exc)}) + + # ------------------------------------------------------------------ + # Flattening helpers + # ------------------------------------------------------------------ + @staticmethod + def flatten_json(data: Any, parent_key: str = "", sep: str = ".") -> Dict[str, Any]: + """Flatten nested dict/list structures using dot-path keys.""" + flat: Dict[str, Any] = {} + + if isinstance(data, dict): + for key, value in data.items(): + new_key = f"{parent_key}{sep}{key}" if parent_key else str(key) + flat.update(DataProcessor.flatten_json(value, new_key, sep)) + return flat + + if isinstance(data, list): + for index, value in enumerate(data): + new_key = f"{parent_key}{sep}{index}" if parent_key else str(index) + flat.update(DataProcessor.flatten_json(value, new_key, sep)) + return flat + + if parent_key: + flat[parent_key] = data + return flat + + @staticmethod + def flatten_dict(data: Dict[str, Any], sep: str = ".") -> Dict[str, Any]: + """Legacy wrapper expected by older callers.""" + return DataProcessor.flatten_json(data, sep=sep) + + # ------------------------------------------------------------------ + # Path-based value access + # ------------------------------------------------------------------ + @staticmethod + def get_value(data: Union[Dict[str, Any], List[Any]], key: str) -> Any: + """Fetch a nested value using a dot-path key.""" + if not key: + return data + + value: Any = data + for part in key.split("."): + if isinstance(value, dict): + value = value.get(part) + elif isinstance(value, list) and part.isdigit(): + index = int(part) + value = value[index] if 0 <= index < len(value) else None + else: + return None + + if value is None: + return None + + return value + + @staticmethod + def set_value(data: Union[Dict[str, Any], List[Any]], key: str, val: Any) -> bool: + """Set a nested value using a dot-path key.""" + if not isinstance(data, (dict, list)) or not key: + return False + + parts = key.split(".") + target: Any = data + + for part in parts[:-1]: + if isinstance(target, dict): + next_target = target.get(part) + if next_target is None: + next_target = {} + target[part] = next_target + target = next_target + elif isinstance(target, list) and part.isdigit(): + index = int(part) + if not (0 <= index < len(target)): + return False + target = target[index] + else: + return False + + last = parts[-1] + if isinstance(target, dict): + target[last] = val + return True + if isinstance(target, list) and last.isdigit(): + index = int(last) + if 0 <= index < len(target): + target[index] = val + return True + return False + + # ------------------------------------------------------------------ + # Text / HTML helpers + # ------------------------------------------------------------------ + @staticmethod + def clean_output_text(text: str) -> str: + """Strip Markdown code fences and unescape literal newlines.""" + if not isinstance(text, str): + raise ValueError("Input must be a string") + + cleaned = text.strip() + if cleaned.startswith("```"): + cleaned = re.sub(r"^```[a-zA-Z0-9_-]*\n?", "", cleaned) + if cleaned.endswith("```"): + cleaned = cleaned[:-3] + + return cleaned.replace("\\n", "\n").strip() + + @staticmethod + def convert_text_to_html(data: Union[str, Dict[str, Any]]) -> str: + """Convert escaped HTML text to rendered HTML text.""" + if isinstance(data, str): + text = data + elif isinstance(data, dict): + if "output_text" not in data: + raise ValueError("Dictionary must contain 'output_text' key for HTML conversion.") + text = data.get("output_text", "") + else: + raise ValueError(f"Invalid input type: {type(data)}. Expected str or dict.") + + if not isinstance(text, str): + raise ValueError("Invalid text for HTML conversion") + + decoded = html.unescape(text) + decoded = re.sub(r">\s+<", "> <", decoded) + return decoded + + # ------------------------------------------------------------------ + # Object conversion helpers + # ------------------------------------------------------------------ + @staticmethod + def convert_object_to_json(obj: Any, _seen: Optional[set] = None) -> Any: + """Recursively coerce arbitrary Python objects into JSON-safe values.""" + return DataProcessor.to_json_compatible(obj, _seen=_seen) diff --git a/zmongo_toolbag/local_vector_search.py b/zmongo_toolbag/local_vector_search.py new file mode 100644 index 0000000..520256a --- /dev/null +++ b/zmongo_toolbag/local_vector_search.py @@ -0,0 +1,201 @@ +""" +LocalVectorSearch – SafeResult-compatible vector similarity search engine +for ZMongo repositories (sync or async). +""" + +from __future__ import annotations + +import asyncio +import logging +import threading +from concurrent.futures import Future +from typing import Any, Dict, List, Optional + +import numpy as np + +from .safe_result import SafeResult + +# Standard imports without leading dots for this environment + +logger = logging.getLogger(__name__) + +def _run_async_in_new_thread(coro) -> Any: + fut: Future[Any] = Future() + def _runner(): + try: + fut.set_result(asyncio.run(coro)) + except Exception as e: + fut.set_exception(e) + t = threading.Thread(target=_runner, daemon=True) + t.start() + return fut.result() + +class LocalVectorSearch: + def __init__( + self, + repository: Any, + collection: str, + embedding_field: str = "embedding", + chunked_embeddings: bool = False, + score_mode: str = "cosine_0_1", + field_key: Optional[str] = None, + vector_key: str = "vectors", + max_docs: int = 1_000_000, + ): + self.repo = repository + self.collection = collection + self.embedding_field = embedding_field + self.chunked_embeddings = chunked_embeddings + self.score_mode = score_mode + self.field_key = field_key + self.vector_key = vector_key + self.max_docs = max_docs + self._index_matrix: Optional[np.ndarray] = None + self._meta_docs: List[Dict[str, Any]] = [] + self._load_lock = asyncio.Lock() + + def clear_index(self) -> None: + """Resets the in-memory index matrix and metadata for graceful shutdown.""" + self._index_matrix = None + self._meta_docs = [] + logger.info("Local vector index cleared.") + + async def _await_repo_result(self, maybe_result: Any) -> SafeResult: + if asyncio.iscoroutine(maybe_result): + maybe_result = await maybe_result + return maybe_result if isinstance(maybe_result, SafeResult) else SafeResult.ok(maybe_result) + + async def _find_all_docs(self) -> SafeResult: + try: + # Standardizing on ZMongo find_many + repo_call = self.repo.find_many(self.collection, query={}, limit=self.max_docs) + return await self._await_repo_result(repo_call) + except Exception as e: + return SafeResult.fail(f"find_many failed: {e}") + + def _extract_vectors_from_doc(self, doc: Dict[str, Any]) -> List[List[float]]: + emb_container = doc.get(self.embedding_field) + if not emb_container: return [] + emb_entry = emb_container + if self.field_key is not None and isinstance(emb_container, dict): + emb_entry = emb_container.get(self.field_key) + if not emb_entry: return [] + vec_data = emb_entry + if isinstance(emb_entry, dict) and self.vector_key in emb_entry: + vec_data = emb_entry.get(self.vector_key) + if not vec_data: return [] + if isinstance(vec_data, list) and vec_data and isinstance(vec_data[0], list): + if self.chunked_embeddings: + return [v for v in vec_data if isinstance(v, list) and v] + return [vec_data[0]] if vec_data[0] else [] + if isinstance(vec_data, list): + return [vec_data] if vec_data else [] + return [] + + async def _ensure_index(self) -> SafeResult: + if self._index_matrix is not None: + return SafeResult.ok({"count": int(self._index_matrix.shape[0])}) + async with self._load_lock: + if self._index_matrix is not None: + return SafeResult.ok({"count": int(self._index_matrix.shape[0])}) + return await self.rebuild_index() + + def rebuild_index_sync(self) -> SafeResult: + try: + loop = asyncio.get_running_loop() + if loop.is_running(): return _run_async_in_new_thread(self.rebuild_index()) + except RuntimeError: pass + return asyncio.run(self.rebuild_index()) + + async def search(self, query_vector: Any, top_k: int = 5) -> SafeResult: + ensure_res = await self._ensure_index() + if not ensure_res.success: return ensure_res + + if self._index_matrix is None or self._index_matrix.size == 0: + return SafeResult.ok([]) + + try: + # FORCE conversion to a 1D float array regardless of input type + q = np.array(query_vector, dtype=float).flatten() + + dim = int(self._index_matrix.shape[1]) + if q.shape[0] != dim: + return SafeResult.fail(f"Dimension mismatch: Query {q.shape[0]} != Index {dim}") + M = self._index_matrix + qn, Mn = np.linalg.norm(q), np.linalg.norm(M, axis=1) + denom = Mn * qn + valid = denom > 0 + cos = np.zeros(M.shape[0], dtype=float) + cos[valid] = (M[valid] @ q) / denom[valid] + scores = self._to_output_scores(cos) + k = min(max(1, int(top_k)), int(M.shape[0])) + idx = np.argpartition(scores, -k)[-k:] + idx = idx[np.argsort(scores[idx])[::-1]] + out = [{"document": self._meta_docs[i], "retrieval_score": float(scores[i])} for i in idx] + return SafeResult.ok(out) + except Exception as e: return SafeResult.fail(str(e)) + + def _to_output_scores(self, cos: np.ndarray) -> np.ndarray: + mode = (self.score_mode or "cosine_0_1").lower().strip() + if mode == "cosine_0_1": return 0.5 * (cos + 1.0) + return cos.astype(float) + + async def rebuild_index(self) -> SafeResult: + """ + Fetches all documents from the repository, extracts vectors, + and builds the in-memory NumPy search matrix. + """ + try: + # 1. Fetch documents from the ZMongo repository + docs_res = await self._find_all_docs() + if not docs_res.success: + return docs_res + + all_docs = docs_res.data if isinstance(docs_res.data, list) else [] + + temp_vectors = [] + temp_meta = [] + expected_dim = None + + # 2. Extract vectors and maintain metadata mapping + for doc in all_docs: + extracted_vectors = self._extract_vectors_from_doc(doc) + for vec in extracted_vectors: + # Validate consistency of dimensions + current_dim = len(vec) + if expected_dim is None: + expected_dim = current_dim + elif current_dim != expected_dim: + logger.warning( + f"Skipping vector with mismatched dimension: {current_dim} (expected {expected_dim})") + continue + + temp_vectors.append(vec) + temp_meta.append(doc) + + # 3. Handle empty datasets gracefully + if not temp_vectors: + self._index_matrix = np.empty((0, 0), dtype=float) + self._meta_docs = [] + logger.info("Rebuild finished: No vectors found.") + return SafeResult.ok({"count": 0, "dim": 0}) + + # 4. Convert to NumPy matrix for fast similarity math + matrix = np.array(temp_vectors, dtype=float) + + # Final safety check for malformed numpy arrays (e.g., ragged nested sequences) + if matrix.dtype == object: + return SafeResult.fail("Failed to build index: Vectors have inconsistent lengths.") + + self._index_matrix = matrix + self._meta_docs = temp_meta + + logger.info(f"Index rebuilt successfully: {matrix.shape[0]} vectors, {matrix.shape[1]} dimensions.") + return SafeResult.ok({ + "count": int(matrix.shape[0]), + "dim": int(matrix.shape[1]) + }) + + except Exception as e: + logger.exception("Failed to rebuild local vector index") + return SafeResult.fail(f"Rebuild index error: {str(e)}") diff --git a/zmongo_toolbag/safe_result.py b/zmongo_toolbag/safe_result.py new file mode 100644 index 0000000..c5029db --- /dev/null +++ b/zmongo_toolbag/safe_result.py @@ -0,0 +1,167 @@ +import json +import logging +from typing import Any, Dict, Optional, Union + +from quart import Response, jsonify + +from .data_processor import DataProcessor + +logger = logging.getLogger(__name__) + + +class SafeResult: + """Unified, BSON-safe, Quart-compatible response wrapper.""" + + def __init__( + self, + success: bool, + data: Optional[Any] = None, + message: Optional[str] = None, + error: Optional[Union[str, Dict[str, Any]]] = None, + status_code: Optional[int] = None, + _raw_data: Optional[Any] = None, + ): + self.success = success + self._raw_data = _raw_data if _raw_data is not None else data + self.data = DataProcessor.to_json_compatible(data) + self.message = message or ("Success" if success else "Error") + self.error = DataProcessor.to_json_compatible(error) + self.status_code = status_code or (200 if success else 400) + + logger.debug( + "[SafeResult.__init__] success=%s, status=%s, message=%s", + self.success, + self.status_code, + self.message, + ) + + @classmethod + def ok( + cls, + data: Optional[Any] = None, + message: Optional[str] = "OK", + status_code: int = 200, + ) -> "SafeResult": + return cls( + True, + data=data, + message=message, + status_code=status_code, + _raw_data=data, + ) + + @classmethod + def fail( + cls, + error: Optional[Union[str, Dict[str, Any], Exception]] = "Error", + data: Optional[Any] = None, + status_code: int = 400, + message: Optional[str] = None, + ) -> "SafeResult": + err_val = str(error) if isinstance(error, Exception) else error + msg_val = message or (str(err_val) if isinstance(err_val, str) else "Error") + + return cls( + False, + data=data, + message=msg_val, + error=err_val, + status_code=status_code, + _raw_data=data, + ) + + def original(self) -> Any: + return self._apply_keymap(self._raw_data) + + def _apply_keymap(self, obj: Any) -> Any: + if isinstance(obj, list): + return [self._apply_keymap(item) for item in obj] + + if isinstance(obj, dict): + keymap = obj.get("__keymap") + new_obj: Dict[str, Any] = {} + for k, v in obj.items(): + if k == "__keymap": + continue + val = self._apply_keymap(v) + if keymap and k in keymap: + new_obj[keymap[k]] = val + else: + new_obj[k] = val + return new_obj + + return obj + + def to_dict(self) -> Dict[str, Any]: + return { + "success": self.success, + "data": self.data, + "message": self.message, + "error": self.error, + "status_code": self.status_code, + } + + def model_dump(self) -> Dict[str, Any]: + return self.to_dict() + + def to_json(self) -> str: + return json.dumps(self.to_dict(), default=str) + + def __repr__(self) -> str: + return ( + f"SafeResult(success={self.success}, data={self.data}, " + f"message={self.message}, error={self.error})" + ) + + def __bool__(self) -> bool: + return self.success + + def to_response(self, override_status: Optional[int] = None): + code = override_status or self.status_code + return jsonify(self.to_dict()), code + + @classmethod + async def from_quart_response(cls, resp: Union[Response, tuple]) -> "SafeResult": + try: + if isinstance(resp, tuple): + resp_obj, code = resp + else: + resp_obj = resp + code = getattr(resp, "status_code", 200) + + body = await resp_obj.get_data(as_text=True) + parsed = json.loads(body) if body else {} + + if isinstance(parsed, dict) and "success" in parsed: + return cls( + success=parsed.get("success", False), + data=parsed.get("data"), + message=parsed.get("message"), + error=parsed.get("error"), + status_code=parsed.get("status_code", code), + ) + + return cls( + success=(200 <= code < 300), + data=parsed, + message=f"Parsed response ({code})", + status_code=code, + ) + + except Exception as e: + return cls.fail(str(e), status_code=500, message="Response parsing failed") + + def log(self, prefix: str = "") -> "SafeResult": + msg = f"{prefix} [{self.status_code}] {self.message}" + if self.success: + logger.info("[SafeResult.log] βœ… %s", msg) + else: + logger.error("[SafeResult.log] ❌ %s | %s", msg, self.error) + return self + + @classmethod + def ensure(cls, condition: bool, message: str = "Condition failed", **kwargs) -> "SafeResult": + if condition: + success_message = kwargs.pop("message", "OK") + return cls.ok(message=success_message, **kwargs) + return cls.fail(error=message, **kwargs) diff --git a/zmongo_toolbag/zembedder.py b/zmongo_toolbag/zembedder.py new file mode 100644 index 0000000..0ef9055 --- /dev/null +++ b/zmongo_toolbag/zembedder.py @@ -0,0 +1,214 @@ +""" +ZEmbedder – Local BGE-M3 Hybrid Embedding manager for Blackwell RTX 5080. +Integrates with ZMongo, LocalVectorSearch, and SafeResult. +""" + +import asyncio +import logging +import os +from pathlib import Path +from typing import Any, Dict, List, Optional, Union + +from dotenv import load_dotenv +from pymilvus.model.hybrid import BGEM3EmbeddingFunction + +# Internal toolbag imports +from .data_processor import DataProcessor +from .safe_result import SafeResult +from .zmongo import ZMongo +from .local_vector_search import LocalVectorSearch + +# Load environment resources +RESOURCE_PATH = Path.home() / ".resources" +load_dotenv(RESOURCE_PATH / ".env") +load_dotenv(RESOURCE_PATH / ".secrets") + +logger = logging.getLogger(__name__) + +# Constants for BGE-M3 +DEFAULT_MODEL_PATH = "/home/comfyuser/.resources/preserved_models/encoders/bge-m3" +DEFAULT_DB_NAME = "wiki_kb" +class ZEmbedder: + """ + Local BGE-M3 Embedding manager. + Produces Dense and Sparse vectors for hybrid RAG pipelines. + """ + def __init__( + self, + db_name: str = DEFAULT_DB_NAME, + model_path: str = DEFAULT_MODEL_PATH, + device: str = "cuda:0", + output_dimensionality: int = 1024 # BGE-M3 default dense dim + ): + self.repository = ZMongo(db_name=db_name) + self.device = device + self.default_output_dim = output_dimensionality + + # Verify local model path + if not os.path.exists(model_path): + logger.error(f"❌ Model not found at {model_path}. Please download it first.") + raise FileNotFoundError(f"BGE-M3 not found at {model_path}") + + # Initialize the Hybrid Model on Blackwell GPU + self.model = BGEM3EmbeddingFunction( + model_name=model_path, + device=device, + use_fp16=True # Optimized for RTX 5080 + ) + + # Internal searcher for backward compatibility + self.vector_search = LocalVectorSearch( + repository=self.repository, + collection="knowledge_base", + embedding_field="embedding.dense", + vector_key="vectors" + ) + logger.info(f"βœ… ZEmbedder (BGE-M3) initialized on {device}") + + @staticmethod + def _to_dense_list(vec: Any) -> List[float]: + if hasattr(vec, "tolist"): + return vec.tolist() + return list(vec) + + @staticmethod + def _to_sparse_dict(sparse: Any) -> Dict[str, float]: + """ + Normalize BGE-M3 sparse output into a plain dict[str, float]. + Handles dict-like outputs and scipy sparse-style objects. + """ + if sparse is None: + return {} + + if isinstance(sparse, dict): + return {str(k): float(v) for k, v in sparse.items()} + + # scipy sparse row / matrix style + if hasattr(sparse, "tocoo"): + coo = sparse.tocoo() + return {str(int(col)): float(val) for col, val in zip(coo.col, coo.data)} + + # some libs may expose indices/data separately + indices = getattr(sparse, "indices", None) + data = getattr(sparse, "data", None) + if indices is not None and data is not None: + return {str(int(i)): float(v) for i, v in zip(indices, data)} + + raise TypeError(f"Unsupported sparse vector type: {type(sparse)!r}") + + async def embed_many(self, texts: List[str]) -> SafeResult: + """Embeds multiple texts using hybrid logic.""" + try: + if not texts: + return SafeResult.ok({"dense": [], "sparse": [], "count": 0}) + + loop = asyncio.get_running_loop() + embeddings = await loop.run_in_executor( + None, + self.model.encode_documents, + texts, + ) + + dense_vectors = [self._to_dense_list(vec) for vec in embeddings["dense"]] + sparse_vectors = [self._to_sparse_dict(vec) for vec in embeddings["sparse"]] + + return SafeResult.ok({ + "dense": dense_vectors, + "sparse": sparse_vectors, + "count": len(texts), + }) + except Exception as e: + logger.exception("Batch embedding failed") + return SafeResult.fail(e) + + async def get_embedding( + self, + text: Optional[str] = None, + *, + collection: Optional[str] = None, + document_id: Optional[Any] = None, + text_field: str = "text", + embedding_field: str = "embedding", + **kwargs, + ) -> SafeResult: + """ + Fetches or embeds text. Returns a SafeResult containing + BSON-safe dense and sparse vectors. + """ + try: + if not text and collection and document_id is not None: + fetch_res = await self.repository.find_one_async( + collection, + {"_id": document_id}, + ) + if not fetch_res.success: + return fetch_res + + text = DataProcessor.get_value(fetch_res.data, text_field) + + if not text: + return SafeResult.fail("No text content provided or found.") + + loop = asyncio.get_running_loop() + embs = await loop.run_in_executor( + None, + self.model.encode_queries, + [text], + ) + + dense_vec = self._to_dense_list(embs["dense"][0]) + sparse_vec = self._to_sparse_dict(embs["sparse"][0]) + + if collection and document_id is not None: + update_payload = { + text_field: text, + embedding_field: { + "dense": dense_vec, + "sparse": sparse_vec, + "model": "bge-m3-local", + "dimensionality": len(dense_vec), + }, + } + + persist_res = await self.repository.insert_or_update( + collection, + {"_id": document_id}, + update_payload, + ) + if not persist_res.success: + return persist_res + + return SafeResult.ok({ + "dense": dense_vec, + "sparse": sparse_vec, + "text": text, + "dimensionality": len(dense_vec), + }) + except Exception as e: + logger.exception("Single embedding generation failed") + return SafeResult.fail(e) + + def close(self): + """Cleanup resources.""" + if hasattr(self, "repository"): + self.repository.close() + if hasattr(self, "vector_search"): + self.vector_search.clear_index() + + +# --- Async Demo --- +if __name__ == "__main__": + async def run_demo(): + embedder = ZEmbedder() + + test_text = "Legal compliance for AI Blackwell architecture." + res = await embedder.get_embedding(text=test_text) + + if res.success: + print(f"βœ… Successfully generated hybrid embeddings.") + print(f"Dense Dim: {res.data['dimensionality']}") + print(f"Sparse Keys: {list(res.data['sparse'].keys())[:5]}...") + + embedder.close() + + asyncio.run(run_demo()) \ No newline at end of file diff --git a/zmongo_toolbag/zmongo.py b/zmongo_toolbag/zmongo.py new file mode 100644 index 0000000..a4740de --- /dev/null +++ b/zmongo_toolbag/zmongo.py @@ -0,0 +1,833 @@ +import asyncio +import json +import logging +import os +import threading +import time +import weakref +from datetime import datetime, timezone +from pathlib import Path +from typing import Any, Dict, List, Optional, Tuple, Union + +from bson import ObjectId +from dotenv import load_dotenv +from motor import motor_asyncio +from pymongo.errors import PyMongoError + +from .data_processor import DataProcessor +from .safe_result import SafeResult + +ENV_PATH = Path.home() / ".resources" / ".env" +load_dotenv(ENV_PATH) + +logger = logging.getLogger(__name__) + + +class ZMongo: + """Production-oriented MongoDB helper with SafeResult-based responses.""" + + DEFAULT_TIMEOUT_SECONDS = 30.0 + + def __init__( + self, + uri: Optional[str] = None, + db_name: Optional[str] = None, + *, + cache_enabled: bool = True, + cache_ttl_seconds: int = 5, + run_sync_timeout_seconds: float = DEFAULT_TIMEOUT_SECONDS, + ) -> None: + self.uri = uri or os.getenv("MONGO_URI", "mongodb://127.0.0.1:27017") + self.db_name = db_name or os.getenv("MONGO_DATABASE_NAME", "test") + + self.cache_enabled = cache_enabled + self.cache_ttl_seconds = max(0, int(cache_ttl_seconds)) + self.run_sync_timeout_seconds = float(run_sync_timeout_seconds) + + self._loop = asyncio.new_event_loop() + self._loop_ready = threading.Event() + self._closed = False + self._close_lock = threading.Lock() + + self._thread = threading.Thread( + target=self._run_event_loop, + name=f"ZMongoLoop-{id(self)}", + daemon=True, + ) + self._thread.start() + self._loop_ready.wait(timeout=5) + + self._async_clients: "weakref.WeakKeyDictionary[asyncio.AbstractEventLoop, motor_asyncio.AsyncIOMotorClient]" = ( + weakref.WeakKeyDictionary() + ) + + self._cache: Dict[str, Dict[str, Tuple[Any, float]]] = {} + self._cache_lock = threading.RLock() + + def _run_event_loop(self) -> None: + asyncio.set_event_loop(self._loop) + self._loop_ready.set() + try: + self._loop.run_forever() + except Exception: + logger.exception("ZMongo background event loop crashed") + + def _ensure_not_closed(self) -> None: + if self._closed: + raise RuntimeError("ZMongo instance is closed") + + def close(self) -> None: + with self._close_lock: + if self._closed: + return + self._closed = True + + try: + for client in list(self._async_clients.values()): + try: + client.close() + except Exception: + logger.debug("Failed closing Mongo client", exc_info=True) + except Exception: + logger.debug("Failed iterating Mongo clients during close", exc_info=True) + + with self._cache_lock: + self._cache.clear() + + try: + if self._loop.is_running(): + self._loop.call_soon_threadsafe(self._loop.stop) + except Exception: + logger.debug("Failed stopping event loop", exc_info=True) + + try: + if self._thread.is_alive(): + self._thread.join(timeout=2.0) + except Exception: + logger.debug("Failed joining background thread", exc_info=True) + + try: + if not self._loop.is_closed(): + self._loop.close() + except Exception: + logger.debug("Failed closing event loop", exc_info=True) + + def __enter__(self) -> "ZMongo": + self._ensure_not_closed() + return self + + def __exit__(self, exc_type, exc, tb) -> None: + self.close() + + async def __aenter__(self) -> "ZMongo": + self._ensure_not_closed() + return self + + async def __aexit__(self, exc_type, exc, tb) -> None: + self.close() + + def _client_for_current_loop(self) -> motor_asyncio.AsyncIOMotorClient: + self._ensure_not_closed() + + try: + loop = asyncio.get_running_loop() + except RuntimeError: + loop = self._loop + + client = self._async_clients.get(loop) + if client is None: + client = motor_asyncio.AsyncIOMotorClient(self.uri) + self._async_clients[loop] = client + return client + + @property + def db(self): + return self._client_for_current_loop()[self.db_name] + + @staticmethod + def _exception_payload(exc: Exception, operation: Optional[str] = None) -> Dict[str, Any]: + payload = { + "error_type": exc.__class__.__name__, + "error": str(exc), + } + if operation: + payload["operation"] = operation + return payload + + @classmethod + def _fail( + cls, + exc: Exception, + *, + operation: Optional[str] = None, + status_code: int = 500, + message: Optional[str] = None, + data: Optional[Any] = None, + ) -> SafeResult: + payload = cls._exception_payload(exc, operation=operation) + return SafeResult.fail( + error=payload, + data=data, + status_code=status_code, + message=message or str(exc), + ) + + def _normalize_query(self, query: Optional[Dict[str, Any]]) -> Dict[str, Any]: + query = dict(query or {}) + if "_id" in query and isinstance(query["_id"], str) and ObjectId.is_valid(query["_id"]): + query["_id"] = ObjectId(query["_id"]) + return query + + @staticmethod + def _json_key(obj: Any) -> str: + safe_obj = DataProcessor.to_json_compatible(obj) + return json.dumps(safe_obj, sort_keys=True, default=str, separators=(",", ":")) + + def _make_cache_key( + self, + *, + operation: str, + coll: str, + query: Optional[Dict[str, Any]] = None, + extra: Optional[Dict[str, Any]] = None, + ) -> str: + payload = { + "operation": operation, + "collection": coll, + "query": self._normalize_query(query), + "extra": extra or {}, + } + return self._json_key(payload) + + def _get_cached(self, coll: str, key: str) -> Optional[Any]: + if not self.cache_enabled or self.cache_ttl_seconds <= 0: + return None + + now = time.time() + with self._cache_lock: + coll_cache = self._cache.get(coll) + if not coll_cache: + return None + + entry = coll_cache.get(key) + if entry is None: + return None + + value, expires_at = entry + if now >= expires_at: + coll_cache.pop(key, None) + if not coll_cache: + self._cache.pop(coll, None) + return None + return value + + def _set_cached(self, coll: str, key: str, value: Any) -> None: + if not self.cache_enabled or self.cache_ttl_seconds <= 0: + return + + expires_at = time.time() + self.cache_ttl_seconds + with self._cache_lock: + self._cache.setdefault(coll, {})[key] = (value, expires_at) + + def clear_cache(self, coll: Optional[str] = None) -> None: + with self._cache_lock: + if coll is None: + self._cache.clear() + else: + self._cache.pop(coll, None) + + def run_sync(self, coro_or_factory, *args, timeout: Optional[float] = None, **kwargs) -> SafeResult: + self._ensure_not_closed() + + if asyncio.iscoroutine(coro_or_factory): + coro = coro_or_factory + else: + coro = coro_or_factory(*args, **kwargs) + + if not asyncio.iscoroutine(coro): + return SafeResult.fail( + error={ + "error_type": "TypeError", + "error": "run_sync expected a coroutine or coroutine factory", + }, + status_code=500, + message="run_sync misuse", + ) + + future = asyncio.run_coroutine_threadsafe(coro, self._loop) + + try: + result = future.result(timeout=timeout or self.run_sync_timeout_seconds) + if isinstance(result, SafeResult): + return result + return SafeResult.ok(result) + except asyncio.TimeoutError as exc: + future.cancel() + return self._fail(exc, operation="run_sync", status_code=504, message="Mongo operation timed out") + except Exception as exc: + future.cancel() + return self._fail(exc, operation="run_sync", status_code=500) + + async def find_one_async( + self, + coll: str, + query: Dict[str, Any], + *, + cache: bool = False, + **kwargs, + ) -> SafeResult: + try: + self._ensure_not_closed() + normalized_query = self._normalize_query(query) + + cache_key = self._make_cache_key( + operation="find_one", + coll=coll, + query=normalized_query, + extra={"kwargs": kwargs}, + ) + + if cache: + cached = self._get_cached(coll, cache_key) + if cached is not None: + return SafeResult.ok( + { + "document": cached, + "cache_hit": True, + "collection": coll, + "query": normalized_query, + } + ) + + doc = await self.db[coll].find_one(normalized_query, **kwargs) + + if cache and doc is not None: + self._set_cached(coll, cache_key, doc) + + return SafeResult.ok( + { + "document": doc, + "cache_hit": False, + "collection": coll, + "query": normalized_query, + } + ) + except Exception as exc: + return self._fail(exc, operation="find_one") + + async def find_many_async( + self, + coll: str, + query: Optional[Dict[str, Any]] = None, + *, + sort: Optional[Union[List[Tuple[str, int]], Tuple[str, int]]] = None, + limit: Optional[int] = 1000, + cache: bool = False, + ) -> SafeResult: + try: + self._ensure_not_closed() + normalized_query = self._normalize_query(query or {}) + + cache_key = self._make_cache_key( + operation="find_many", + coll=coll, + query=normalized_query, + extra={"sort": sort, "limit": limit}, + ) + + if cache: + cached = self._get_cached(coll, cache_key) + if cached is not None: + return SafeResult.ok( + { + "documents": cached, + "count": len(cached), + "cache_hit": True, + "collection": coll, + "query": normalized_query, + } + ) + + cursor = self.db[coll].find(normalized_query) + if sort: + cursor = cursor.sort(sort) + if limit is not None: + cursor = cursor.limit(limit) + + docs = await cursor.to_list(length=limit) + if cache: + self._set_cached(coll, cache_key, docs) + + return SafeResult.ok( + { + "documents": docs, + "count": len(docs), + "cache_hit": False, + "collection": coll, + "query": normalized_query, + } + ) + except Exception as exc: + return self._fail(exc, operation="find_many") + + async def aggregate_async(self, coll: str, pipeline: List[Dict[str, Any]]) -> SafeResult: + try: + self._ensure_not_closed() + cursor = self.db[coll].aggregate(pipeline) + docs = await cursor.to_list(length=None) + return SafeResult.ok( + { + "documents": docs, + "count": len(docs), + "collection": coll, + "pipeline": pipeline, + } + ) + except Exception as exc: + return self._fail(exc, operation="aggregate") + + async def count_documents_async( + self, + coll: str, + query: Optional[Dict[str, Any]] = None, + *, + cache: bool = False, + ) -> SafeResult: + try: + self._ensure_not_closed() + normalized_query = self._normalize_query(query or {}) + + cache_key = self._make_cache_key( + operation="count_documents", + coll=coll, + query=normalized_query, + ) + + if cache: + cached = self._get_cached(coll, cache_key) + if cached is not None: + return SafeResult.ok( + { + "count": cached, + "cache_hit": True, + "collection": coll, + "query": normalized_query, + } + ) + + count = await self.db[coll].count_documents(normalized_query) + if cache: + self._set_cached(coll, cache_key, count) + + return SafeResult.ok( + { + "count": count, + "cache_hit": False, + "collection": coll, + "query": normalized_query, + } + ) + except Exception as exc: + return self._fail(exc, operation="count_documents") + + async def insert_one_async(self, coll: str, doc: Dict[str, Any]) -> SafeResult: + try: + self._ensure_not_closed() + result = await self.db[coll].insert_one(doc) + self.clear_cache(coll) + return SafeResult.ok({"inserted_id": result.inserted_id, "collection": coll}) + except Exception as exc: + return self._fail(exc, operation="insert_one") + + async def insert_many_async(self, coll: str, docs: List[Dict[str, Any]]) -> SafeResult: + try: + self._ensure_not_closed() + result = await self.db[coll].insert_many(docs) + self.clear_cache(coll) + return SafeResult.ok( + { + "inserted_ids": result.inserted_ids, + "inserted_count": len(result.inserted_ids), + "collection": coll, + } + ) + except Exception as exc: + return self._fail(exc, operation="insert_many") + + async def update_one_async( + self, + coll: str, + query: Dict[str, Any], + update: Dict[str, Any], + *, + upsert: bool = False, + ) -> SafeResult: + try: + self._ensure_not_closed() + normalized_query = self._normalize_query(query) + update_doc = update if any(k.startswith("$") for k in update.keys()) else {"$set": update} + + result = await self.db[coll].update_one(normalized_query, update_doc, upsert=upsert) + self.clear_cache(coll) + + return SafeResult.ok( + { + "matched_count": result.matched_count, + "modified_count": result.modified_count, + "upserted_id": result.upserted_id, + "collection": coll, + "query": normalized_query, + } + ) + except Exception as exc: + return self._fail(exc, operation="update_one") + + async def update_many_async( + self, + coll: str, + query_or_ops: Union[Dict[str, Any], List[Any]], + update: Optional[Dict[str, Any]] = None, + *, + upsert: bool = False, + ) -> SafeResult: + try: + self._ensure_not_closed() + + if isinstance(query_or_ops, list): + result = await self.db[coll].bulk_write(query_or_ops) + self.clear_cache(coll) + return SafeResult.ok( + { + "bulk_api_result": getattr(result, "bulk_api_result", None), + "inserted_count": getattr(result, "inserted_count", None), + "matched_count": getattr(result, "matched_count", None), + "modified_count": getattr(result, "modified_count", None), + "deleted_count": getattr(result, "deleted_count", None), + "upserted_count": getattr(result, "upserted_count", None), + "collection": coll, + } + ) + + normalized_query = self._normalize_query(query_or_ops) + if update is None: + return SafeResult.fail( + error={"error_type": "ValueError", "error": "update is required for non-bulk update_many"}, + status_code=400, + message="Missing update document", + ) + + update_doc = update if any(k.startswith("$") for k in update.keys()) else {"$set": update} + result = await self.db[coll].update_many(normalized_query, update_doc, upsert=upsert) + self.clear_cache(coll) + + return SafeResult.ok( + { + "matched_count": result.matched_count, + "modified_count": result.modified_count, + "collection": coll, + "query": normalized_query, + } + ) + except Exception as exc: + return self._fail(exc, operation="update_many") + + async def delete_one_async(self, coll: str, query: Dict[str, Any]) -> SafeResult: + try: + self._ensure_not_closed() + normalized_query = self._normalize_query(query) + result = await self.db[coll].delete_one(normalized_query) + self.clear_cache(coll) + return SafeResult.ok( + {"deleted_count": result.deleted_count, "collection": coll, "query": normalized_query} + ) + except Exception as exc: + return self._fail(exc, operation="delete_one") + + async def delete_many_async(self, coll: str, query: Dict[str, Any]) -> SafeResult: + try: + self._ensure_not_closed() + normalized_query = self._normalize_query(query) + result = await self.db[coll].delete_many(normalized_query) + self.clear_cache(coll) + return SafeResult.ok( + {"deleted_count": result.deleted_count, "collection": coll, "query": normalized_query} + ) + except Exception as exc: + return self._fail(exc, operation="delete_many") + + async def insert_or_update_async( + self, + coll: str, + query_or_document: Dict[str, Any], + data: Optional[Dict[str, Any]] = None, + ) -> SafeResult: + try: + self._ensure_not_closed() + + if data is None: + if "_id" not in query_or_document: + return SafeResult.fail( + error={"error_type": "ValueError", "error": "Document must include _id when data is omitted"}, + status_code=400, + message="Missing _id for insert_or_update", + ) + + query = self._normalize_query({"_id": query_or_document["_id"]}) + update_doc = {"$set": {k: v for k, v in query_or_document.items() if k != "_id"}} + else: + query = self._normalize_query(query_or_document) + update_doc = data if any(k.startswith("$") for k in data.keys()) else {"$set": data} + + result = await self.db[coll].update_one(query, update_doc, upsert=True) + self.clear_cache(coll) + + return SafeResult.ok( + { + "matched_count": result.matched_count, + "modified_count": result.modified_count, + "upserted_id": result.upserted_id, + "collection": coll, + "query": query, + } + ) + except Exception as exc: + return self._fail(exc, operation="insert_or_update") + + async def save_value_async( + self, + coll: str, + value: Any, + *, + query: Optional[Dict[str, Any]] = None, + field_path: Optional[str] = None, + upsert: bool = True, + parse_json_strings: bool = True, + normalize_for_storage: bool = False, + ) -> SafeResult: + try: + self._ensure_not_closed() + + normalized_query = self._normalize_query(query or {}) + parsed_value = value + + if parse_json_strings and isinstance(value, str): + stripped = value.strip() + if stripped: + try: + parsed_value = json.loads(stripped) + except Exception: + parsed_value = value + + if normalize_for_storage: + parsed_value = DataProcessor.to_json_compatible(parsed_value) + + if not normalized_query and isinstance(parsed_value, dict) and parsed_value.get("_id") is not None: + normalized_query = self._normalize_query({"_id": parsed_value["_id"]}) + + if field_path and str(field_path).strip(): + update_payload = {str(field_path).strip(): parsed_value} + elif isinstance(parsed_value, dict): + update_payload = dict(parsed_value) + update_payload.pop("_id", None) + else: + update_payload = {"value": parsed_value} + + if not normalized_query and not upsert: + return SafeResult.fail( + error={"error_type": "ValueError", "error": "No query provided and upsert is False"}, + status_code=400, + message="Cannot determine target document", + ) + + if not normalized_query: + result = await self.db[coll].insert_one(update_payload) + self.clear_cache(coll) + return SafeResult.ok( + { + "operation": "inserted_new", + "inserted_id": result.inserted_id, + "collection": coll, + "field_path": field_path, + "saved_value": parsed_value, + } + ) + + result = await self.db[coll].update_one( + normalized_query, + {"$set": update_payload}, + upsert=upsert, + ) + self.clear_cache(coll) + + return SafeResult.ok( + { + "operation": "inserted_via_upsert" if result.upserted_id is not None else "updated_existing", + "matched_count": result.matched_count, + "modified_count": result.modified_count, + "upserted_id": result.upserted_id, + "collection": coll, + "query": normalized_query, + "field_path": field_path, + "saved_value": parsed_value, + } + ) + except Exception as exc: + return self._fail(exc, operation="save_value") + + async def list_collections_async(self) -> SafeResult: + try: + self._ensure_not_closed() + names = await self.db.list_collection_names() + return SafeResult.ok({"collections": names, "count": len(names)}) + except Exception as exc: + return self._fail(exc, operation="list_collections") + + async def drop_database_async(self, database_name: Optional[str] = None) -> SafeResult: + try: + self._ensure_not_closed() + name = database_name or self.db_name + await self._client_for_current_loop().drop_database(name) + self.clear_cache() + return SafeResult.ok({"dropped_database": name}) + except Exception as exc: + return self._fail(exc, operation="drop_database") + + async def ping_async(self) -> SafeResult: + try: + self._ensure_not_closed() + start = time.time() + await self.db.command({"ping": 1}) + end = time.time() + return SafeResult.ok({"ok": True, "latency_seconds": end - start, "database": self.db_name}) + except Exception as exc: + return self._fail(exc, operation="ping") + + async def sync_timestamp_async(self) -> SafeResult: + try: + self._ensure_not_closed() + start = time.time() + + try: + result = await self.db.command({"hello": 1}) + except PyMongoError: + result = await self.db.command({"isMaster": 1}) + + end = time.time() + server_time = result.get("localTime") + if server_time and server_time.tzinfo is None: + server_time = server_time.replace(tzinfo=timezone.utc) + + return SafeResult.ok( + { + "server_time": server_time, + "latency_seconds": end - start, + "offset_seconds": ( + (server_time - datetime.now(timezone.utc)).total_seconds() if server_time is not None else None + ), + } + ) + except Exception as exc: + return self._fail(exc, operation="sync_timestamp") + + def find_one(self, coll: str, query: Dict[str, Any], *, cache: bool = False, **kwargs) -> SafeResult: + return self.run_sync(self.find_one_async, coll, query, cache=cache, **kwargs) + + def find_many( + self, + coll: str, + query: Optional[Dict[str, Any]] = None, + *, + sort: Optional[Union[List[Tuple[str, int]], Tuple[str, int]]] = None, + limit: Optional[int] = 1000, + cache: bool = False, + ) -> SafeResult: + return self.run_sync(self.find_many_async, coll, query, sort=sort, limit=limit, cache=cache) + + def aggregate(self, coll: str, pipeline: List[Dict[str, Any]]) -> SafeResult: + return self.run_sync(self.aggregate_async, coll, pipeline) + + def count_documents(self, coll: str, query: Optional[Dict[str, Any]] = None, *, cache: bool = False) -> SafeResult: + return self.run_sync(self.count_documents_async, coll, query, cache=cache) + + def insert_one(self, coll: str, doc: Dict[str, Any]) -> SafeResult: + return self.run_sync(self.insert_one_async, coll, doc) + + def insert_many(self, coll: str, docs: List[Dict[str, Any]]) -> SafeResult: + return self.run_sync(self.insert_many_async, coll, docs) + + def update_one( + self, + coll: str, + query: Dict[str, Any], + update: Dict[str, Any], + *, + upsert: bool = False, + ) -> SafeResult: + return self.run_sync(self.update_one_async, coll, query, update, upsert=upsert) + + def update_many( + self, + coll: str, + query_or_ops: Union[Dict[str, Any], List[Any]], + update: Optional[Dict[str, Any]] = None, + *, + upsert: bool = False, + ) -> SafeResult: + return self.run_sync(self.update_many_async, coll, query_or_ops, update, upsert=upsert) + + def delete_one(self, coll: str, query: Dict[str, Any]) -> SafeResult: + return self.run_sync(self.delete_one_async, coll, query) + + def delete_many(self, coll: str, query: Dict[str, Any]) -> SafeResult: + return self.run_sync(self.delete_many_async, coll, query) + + def delete_all_documents(self, coll: str) -> SafeResult: + return self.run_sync(self.delete_many_async, coll, {}) + + def delete_documents(self, coll: str, query: Dict[str, Any]) -> SafeResult: + return self.run_sync(self.delete_many_async, coll, query) + + def insert_documents(self, coll: str, docs: List[Dict[str, Any]]) -> SafeResult: + return self.run_sync(self.insert_many_async, coll, docs) + + def insert_or_update( + self, + coll: str, + query_or_document: Dict[str, Any], + data: Optional[Dict[str, Any]] = None, + ) -> SafeResult: + return self.run_sync(self.insert_or_update_async, coll, query_or_document, data) + + def save_value( + self, + coll: str, + value: Any, + *, + query: Optional[Dict[str, Any]] = None, + field_path: Optional[str] = None, + upsert: bool = True, + parse_json_strings: bool = True, + normalize_for_storage: bool = False, + ) -> SafeResult: + return self.run_sync( + self.save_value_async, + coll, + value, + query=query, + field_path=field_path, + upsert=upsert, + parse_json_strings=parse_json_strings, + normalize_for_storage=normalize_for_storage, + ) + + def list_collections(self) -> SafeResult: + return self.run_sync(self.list_collections_async) + + def drop_database(self, database_name: Optional[str] = None) -> SafeResult: + return self.run_sync(self.drop_database_async, database_name) + + def ping(self) -> SafeResult: + return self.run_sync(self.ping_async) + + def sync_timestamp(self) -> SafeResult: + return self.run_sync(self.sync_timestamp_async)