initial commit

1 files changed, 190 insertions, 0 deletions

diff --git a/harness.py b/harness.py
new file mode 100644
index 0000000..f588328
--- /dev/null
+++ b/harness.py
@@ -0,0 +1,190 @@
+#!/usr/bin/env python3
+import argparse
+import asyncio
+import json
+import pathlib
+import sys
+import types
+from types import SimpleNamespace
+from typing import Any
+
+# ---------------------------------------------------------------------------
+# 1. Stub out `fastapi` and `open_webui.*` in sys.modules BEFORE importing tool
+# ---------------------------------------------------------------------------
+def install_open_webui_stubs(
+    comfyui_url: str,
+    api_key: str,
+    workflow_path: pathlib.Path,
+    nodes_config: list[dict[str, Any]],
+    image_size: str,
+    image_steps: int,
+    output_dir: pathlib.Path,
+) -> None:
+    # Ensure fastapi.Request exists even in a minimal venv
+    if "fastapi" not in sys.modules:
+        fastapi_mod = types.ModuleType("fastapi")
+        fastapi_mod.Request = object
+        sys.modules["fastapi"] = fastapi_mod
+
+    workflow_raw = workflow_path.read_text(encoding="utf-8")
+
+    # Mock `open_webui.models.config.Config`
+    config_mod = types.ModuleType("open_webui.models.config")
+
+    class MockConfig:
+        @staticmethod
+        async def get(key: str, default: Any = None) -> Any:
+            return default
+
+    config_mod.Config = MockConfig
+
+    # Mock `open_webui.models.users.Users`
+    users_mod = types.ModuleType("open_webui.models.users")
+
+    class MockUsers:
+        @staticmethod
+        def get_user_by_id(user_id: str) -> Any:
+            return SimpleNamespace(id=user_id, name="Felix")
+
+    users_mod.Users = MockUsers
+
+    # Mock `open_webui.routers.images`
+    images_mod = types.ModuleType("open_webui.routers.images")
+
+    async def mock_get_image_config() -> Any:
+        return SimpleNamespace(
+            COMFYUI_BASE_URL=comfyui_url,
+            COMFYUI_API_KEY=api_key,
+            COMFYUI_WORKFLOW=workflow_raw,
+            COMFYUI_WORKFLOW_NODES=nodes_config,
+            IMAGE_SIZE=image_size,
+            IMAGE_STEPS=image_steps,
+        )
+
+    def mock_upload_image(
+        request: Any,
+        image_data: bytes,
+        content_type: str,
+        metadata: dict[str, Any],
+        user: Any,
+    ) -> Any:
+        output_dir.mkdir(parents=True, exist_ok=True)
+        ext = ".png" if "png" in content_type else ".jpg"
+        file_id = f"seed_{metadata.get('seed', 'unknown')}_{int(asyncio.get_event_loop().time() * 1000)}"
+        out_file = output_dir / f"{file_id}{ext}"
+        out_file.write_bytes(image_data)
+        print(f"  [Storage] Saved {len(image_data)} bytes -> {out_file}")
+        # Return a mock FileModel with an `.id` attribute just like Open WebUI
+        return SimpleNamespace(id=file_id)
+
+    images_mod.get_image_config = mock_get_image_config
+    images_mod.upload_image = mock_upload_image
+
+    # Register package hierarchy in sys.modules
+    sys.modules["open_webui"] = types.ModuleType("open_webui")
+    sys.modules["open_webui.models"] = types.ModuleType("open_webui.models")
+    sys.modules["open_webui.routers"] = types.ModuleType("open_webui.routers")
+    sys.modules["open_webui.models.config"] = config_mod
+    sys.modules["open_webui.models.users"] = users_mod
+    sys.modules["open_webui.routers.images"] = images_mod
+
+
+# ---------------------------------------------------------------------------
+# 2. Terminal Event Emitter & Interactive Event Call
+# ---------------------------------------------------------------------------
+async def console_event_emitter(event: dict[str, Any]) -> None:
+    ev_type = event.get("type")
+    data = event.get("data", {})
+    if ev_type == "status":
+        done_marker = "✓" if data.get("done") else "…"
+        print(f"  [{done_marker}] {data.get('description')}")
+    elif ev_type == "files":
+        print(f"  [Files Emitted] {json.dumps(data.get('files', []))}")
+    else:
+        print(f"  [Event:{ev_type}] {data}")
+
+
+async def console_event_call(event: dict[str, Any]) -> bool:
+    ev_type = event.get("type")
+    data = event.get("data", {})
+    print(f"\n--- MODAL ({ev_type}): {data.get('title')} ---")
+    print(data.get("message", ""))
+    loop = asyncio.get_running_loop()
+    answer = await loop.run_in_executor(
+        None, lambda: input("Confirm / Retry? [y/N]: ").strip().lower()
+    )
+    return answer in ("y", "yes")
+
+
+# ---------------------------------------------------------------------------
+# 3. CLI Runner
+# ---------------------------------------------------------------------------
+async def main() -> None:
+    parser = argparse.ArgumentParser(description="Standalone test harness for Comfier")
+    parser.add_argument("--url", default="http://10.42.0.146:8188", help="ComfyUI base URL")
+    parser.add_argument("--workflow", type=pathlib.Path, required=True, help="Path to ComfyUI API workflow JSON")
+    parser.add_argument("--prompt", default="a tiny blue marble on a white surface", help="Positive prompt")
+    parser.add_argument("--neg", default=None, help="Negative prompt")
+    parser.add_argument("--size", default="256x256", help="Image size (WIDTHxHEIGHT)")
+    parser.add_argument("--seed", default="123", help="Seed")
+    parser.add_argument("--steps", type=int, default=30, help="Global IMAGE_STEPS")
+    parser.add_argument("--neg-nodes", default="459:452", help="Valve: negative_prompt_node_ids")
+    parser.add_argument("--cancel-after", type=float, default=None, help="Simulate user clicking Stop after N seconds")
+    args = parser.parse_args()
+
+    # Match your Open WebUI global COMFYUI_WORKFLOW_NODES setup
+    global_nodes = [
+        {"type": "prompt", "key": "value", "node_ids": ["520:506"]},
+        {"type": "seed", "key": "seed", "node_ids": ["520:517"]},
+        {"type": "width", "key": "value", "node_ids": ["520:518"]},
+        {"type": "height", "key": "value", "node_ids": ["520:519"]},
+        {"type": "steps", "key": "steps", "node_ids": ["459:458"]},
+    ]
+
+    install_open_webui_stubs(
+        comfyui_url=args.url,
+        api_key="",
+        workflow_path=args.workflow,
+        nodes_config=global_nodes,
+        image_size="1024x1024",
+        image_steps=args.steps,
+        output_dir=pathlib.Path("./test_out"),
+    )
+
+    # Import your tool AFTER stubs are installed
+    from comfier import Tools
+
+    tools = Tools()
+    tools.valves.negative_prompt_node_ids = args.neg_nodes
+
+    task = asyncio.create_task(
+        tools.comfier_generate_image(
+            positive_prompt=args.prompt,
+            negative_prompt=args.neg,
+            size=args.size,
+            seed=args.seed,
+            __request__=SimpleNamespace(),
+            __user__={"id": "local-dev"},
+            __event_emitter__=console_event_emitter,
+            __event_call__=console_event_call,
+        )
+    )
+
+    if args.cancel_after is not None:
+        await asyncio.sleep(args.cancel_after)
+        print(f"\n[Harness] Simulating user abort after {args.cancel_after}s...")
+        task.cancel()
+
+    try:
+        result = await task
+        print("\n[Tool Return Payload]")
+        print(json.dumps(result, indent=2))
+    except asyncio.CancelledError:
+        print("\n[Harness] Task cleanly cancelled.")
+
+
+if __name__ == "__main__":
+    try:
+        asyncio.run(main())
+    except KeyboardInterrupt:
+        print("\nInterrupted via Ctrl+C.")