Ethanfel
fcca25397c
Record usage in a daemon thread to avoid blocking the event loop with SQLite fsync on slow storage. Replace text button label with a bar-chart-with-nodes SVG icon in the legacy menu fallback. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
74 lines
2.1 KiB
Python
74 lines
2.1 KiB
Python
import logging
|
|
import threading
|
|
|
|
from aiohttp import web
|
|
from server import PromptServer
|
|
|
|
from .mapper import NodePackageMapper
|
|
from .tracker import UsageTracker
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
NODE_CLASS_MAPPINGS = {}
|
|
WEB_DIRECTORY = "./js"
|
|
|
|
mapper = NodePackageMapper()
|
|
tracker = UsageTracker()
|
|
|
|
|
|
def on_prompt_handler(json_data):
|
|
"""Called on every prompt submission. Extracts class_types and records usage."""
|
|
try:
|
|
prompt = json_data.get("prompt", {})
|
|
class_types = set()
|
|
for node_id, node_data in prompt.items():
|
|
ct = node_data.get("class_type")
|
|
if ct:
|
|
class_types.add(ct)
|
|
if class_types:
|
|
threading.Thread(
|
|
target=tracker.record_usage,
|
|
args=(class_types, mapper),
|
|
daemon=True,
|
|
).start()
|
|
except Exception:
|
|
logger.warning("nodes-stats: error recording usage", exc_info=True)
|
|
return json_data
|
|
|
|
|
|
PromptServer.instance.add_on_prompt_handler(on_prompt_handler)
|
|
|
|
|
|
routes = PromptServer.instance.routes
|
|
|
|
|
|
@routes.get("/nodes-stats/packages")
|
|
async def get_package_stats(request):
|
|
try:
|
|
stats = tracker.get_package_stats(mapper)
|
|
return web.json_response(stats)
|
|
except Exception:
|
|
logger.error("nodes-stats: error getting package stats", exc_info=True)
|
|
return web.json_response({"error": "internal error"}, status=500)
|
|
|
|
|
|
@routes.get("/nodes-stats/usage")
|
|
async def get_node_stats(request):
|
|
try:
|
|
stats = tracker.get_node_stats()
|
|
return web.json_response(stats)
|
|
except Exception:
|
|
logger.error("nodes-stats: error getting node stats", exc_info=True)
|
|
return web.json_response({"error": "internal error"}, status=500)
|
|
|
|
|
|
@routes.post("/nodes-stats/reset")
|
|
async def reset_stats(request):
|
|
try:
|
|
tracker.reset()
|
|
mapper.invalidate()
|
|
return web.json_response({"status": "ok"})
|
|
except Exception:
|
|
logger.error("nodes-stats: error resetting stats", exc_info=True)
|
|
return web.json_response({"error": "internal error"}, status=500)
|