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:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user