diff options
Diffstat (limited to 'harness.py')
| -rw-r--r-- | harness.py | 190 |
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.") |