""" REST API routes for Prompt Database CRUD operations. Registered with ComfyUI's PromptServer on import. """ import os from aiohttp import web from server import PromptServer from . import db HTML_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "web") routes = PromptServer.instance.routes @routes.get("/prompt_db/manager") async def serve_manager(request): """Serve the Prompt Database Manager HTML page.""" html_path = os.path.join(HTML_DIR, "prompt_manager.html") return web.FileResponse(html_path) @routes.get("/prompt_db/categories") async def get_categories(request): categories = db.get_categories() return web.json_response({"categories": categories}) @routes.get("/prompt_db/prompts") async def get_prompts(request): category = request.query.get("category", None) prompts = db.get_prompts(category) return web.json_response({"prompts": prompts}) @routes.get("/prompt_db/all") async def get_all(request): data = db.load_db() return web.json_response(data) @routes.get("/prompt_db/loras") async def get_loras(request): """Return list of available LoRA files from ComfyUI's models folder.""" loras = db.get_available_loras() return web.json_response({"loras": loras}) @routes.post("/prompt_db/prompt") async def add_prompt(request): try: body = await request.json() category = body.get("category", "") name = body.get("name", "") text = body.get("text", "") loras = body.get("loras", []) if not category or not name: return web.json_response( {"error": "category and name are required"}, status=400 ) prompt = db.add_prompt(category, name, text, loras=loras) return web.json_response({"prompt": prompt}) except Exception as e: return web.json_response({"error": str(e)}, status=500) @routes.put("/prompt_db/prompt/{prompt_id}") async def update_prompt(request): try: prompt_id = request.match_info["prompt_id"] body = await request.json() result = db.update_prompt( prompt_id, category=body.get("category"), name=body.get("name"), text=body.get("text"), loras=body.get("loras"), ) if result: return web.json_response({"prompt": result}) return web.json_response({"error": "Prompt not found"}, status=404) except Exception as e: return web.json_response({"error": str(e)}, status=500) @routes.delete("/prompt_db/prompt/{prompt_id}") async def delete_prompt(request): try: prompt_id = request.match_info["prompt_id"] if db.delete_prompt(prompt_id): return web.json_response({"success": True}) return web.json_response({"error": "Prompt not found"}, status=404) except Exception as e: return web.json_response({"error": str(e)}, status=500) @routes.post("/prompt_db/category") async def add_category(request): try: body = await request.json() name = body.get("name", "") if not name: return web.json_response({"error": "name is required"}, status=400) if db.add_category(name): return web.json_response({"success": True}) return web.json_response({"error": "Category already exists"}, status=409) except Exception as e: return web.json_response({"error": str(e)}, status=500) @routes.delete("/prompt_db/category/{name}") async def delete_category(request): try: name = request.match_info["name"] if db.delete_category(name): return web.json_response({"success": True}) return web.json_response({"error": "Category not found"}, status=404) except Exception as e: return web.json_response({"error": str(e)}, status=500) @routes.put("/prompt_db/categories/reorder") async def reorder_categories(request): try: body = await request.json() ordered = body.get("categories", []) if not ordered: return web.json_response({"error": "categories list required"}, status=400) if db.reorder_categories(ordered): return web.json_response({"success": True}) return web.json_response( {"error": "Category list does not match existing categories"}, status=400 ) except Exception as e: return web.json_response({"error": str(e)}, status=500) @routes.put("/prompt_db/prompts/reorder") async def reorder_prompts(request): try: body = await request.json() ordered_ids = body.get("prompt_ids", []) if not ordered_ids: return web.json_response({"error": "prompt_ids list required"}, status=400) if db.reorder_prompts(ordered_ids): return web.json_response({"success": True}) return web.json_response({"error": "Reorder failed"}, status=400) except Exception as e: return web.json_response({"error": str(e)}, status=500) @routes.get("/prompt_db/export") async def export_db_route(request): try: data = db.export_db() return web.json_response(data) except Exception as e: return web.json_response({"error": str(e)}, status=500) @routes.post("/prompt_db/import") async def import_db_route(request): try: body = await request.json() mode = body.get("mode", "merge") incoming = body.get("data") if not incoming or "categories" not in incoming or "prompts" not in incoming: return web.json_response( {"error": "Invalid import data: must have categories and prompts"}, status=400, ) stats = db.import_db(incoming, mode) return web.json_response({"success": True, **stats}) except Exception as e: return web.json_response({"error": str(e)}, status=500)