feat: add stats API endpoints and frontend JS with proxy injection

- Add /api/lm-extra/fetch-stats (POST) endpoint that bulk-fetches CivitAI
  stats for all remote models and stores them in the local stats DB with
  optional WebSocket progress broadcasting
- Add /api/lm-extra/stats-status (GET) endpoint to check DB population
- Add /lm-extra/stats-ui.js handler that serves the client-side JS file
- Create static/lm_stats_ui.js: intercepts list API responses to capture
  stats, patches model cards with download/rating/likes badges, adds
  CivitAI Stats sort options to the dropdown, and adds a Fetch Stats
  toolbar button
- Inject stats-ui.js script tag into proxied HTML page responses

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
2026-03-15 14:39:56 +01:00
parent 1ce1e7d48a
commit 467aea0f47
2 changed files with 335 additions and 0 deletions
+112
View File
@@ -13,6 +13,7 @@ from __future__ import annotations
import asyncio
import json
import logging
from pathlib import Path
import aiohttp
from aiohttp import web, WSMsgType
@@ -190,6 +191,8 @@ _SEND_SYNC_HANDLERS = {
# Stats enrichment
# ---------------------------------------------------------------------------
_STATS_JS_PATH = Path(__file__).resolve().parent / "static" / "lm_stats_ui.js"
_stats_db = None
@@ -333,6 +336,93 @@ async def _proxy_list_with_stats_sort(
})
# ---------------------------------------------------------------------------
# Extra API endpoints (stats)
# ---------------------------------------------------------------------------
async def _handle_fetch_stats(request: web.Request) -> web.Response:
"""Bulk fetch CivitAI stats for all models."""
try:
from .stats_service import StatsFetchService
except ImportError:
from stats_service import StatsFetchService
db = _get_stats_db()
# Get the lora list from remote
session = await _get_proxy_session()
remote_url = f"{remote_config.remote_url}/api/lm/loras/list?page=1&page_size=9999"
try:
async with session.get(remote_url) as resp:
if resp.status != 200:
return web.json_response(
{"success": False, "error": "Failed to fetch lora list"},
status=500,
)
data = await resp.json()
except Exception as exc:
return web.json_response(
{"success": False, "error": f"Failed to fetch lora list: {exc}"},
status=500,
)
items = data.get("items", [])
models = []
for item in items:
civitai = item.get("civitai") or {}
model_id = civitai.get("modelId")
sha256 = item.get("sha256")
if model_id and sha256:
models.append({"sha256": sha256, "civitai_model_id": model_id})
if not models:
return web.json_response({"success": True, "updated": 0,
"message": "No models with CivitAI IDs found"})
service = StatsFetchService(db, api_key=remote_config.civitai_api_key or None)
try:
try:
server = _get_prompt_server()
async def progress(current, total):
server.send_sync("lm_stats_progress", {
"current": current, "total": total,
})
updated = await service.fetch_stats_for_models(models, progress_callback=progress)
except Exception:
updated = await service.fetch_stats_for_models(models)
finally:
await service.close()
return web.json_response({"success": True, "updated": updated, "total": len(models)})
async def _handle_stats_status(request: web.Request) -> web.Response:
"""Check if stats DB is populated."""
db = _get_stats_db()
return web.json_response({"success": True, "count": db.count()})
async def _handle_stats_ui_js(request: web.Request) -> web.Response:
"""Serve the stats UI JavaScript file."""
try:
content = _STATS_JS_PATH.read_text(encoding="utf-8")
return web.Response(
text=content,
content_type="application/javascript",
)
except Exception as exc:
return web.Response(status=404, text=str(exc))
_EXTRA_API_HANDLERS = {
"/api/lm-extra/fetch-stats": _handle_fetch_stats,
"/api/lm-extra/stats-status": _handle_stats_status,
"/lm-extra/stats-ui.js": _handle_stats_ui_js,
}
# Shared HTTP session for proxied requests (connection pooling)
_proxy_session: aiohttp.ClientSession | None = None
@@ -485,6 +575,11 @@ async def lm_remote_proxy_middleware(request: web.Request, handler):
if _is_ws_route(path):
return await _proxy_ws(request)
# Extra API endpoints (stats, etc.)
extra_handler = _EXTRA_API_HANDLERS.get(path)
if extra_handler is not None:
return await extra_handler(request)
# Stats-based sorting: intercept list requests, fetch all from remote,
# sort locally, paginate locally
if (_should_proxy(path) and path.endswith("/list")
@@ -506,6 +601,23 @@ async def lm_remote_proxy_middleware(request: web.Request, handler):
return web.json_response(enriched)
except Exception:
pass # Return original response on any error
# Inject stats UI script into HTML pages
if (path in _PROXY_PAGE_ROUTES or path.rstrip("/") in _PROXY_PAGE_ROUTES):
if response.status == 200:
content_type = response.headers.get("Content-Type", "")
if "text/html" in content_type:
try:
html = response.body.decode("utf-8")
inject = '<script src="/lm-extra/stats-ui.js"></script>'
html = html.replace("</body>", f"{inject}\n</body>")
return web.Response(
text=html,
content_type="text/html",
charset="utf-8",
status=response.status,
)
except Exception:
pass
return response
# Not a LoRA Manager route — fall through