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
|
||||
|
||||
@@ -0,0 +1,223 @@
|
||||
// static/lm_stats_ui.js
|
||||
/**
|
||||
* CivitAI Stats UI — card badges, sort dropdown options, fetch button.
|
||||
*
|
||||
* Injected into LoRA Manager pages by the LM-Remote proxy.
|
||||
* Reads stats from enriched /api/lm/loras/list responses and
|
||||
* patches the DOM to show badges, sort options, and a fetch button.
|
||||
*/
|
||||
(function () {
|
||||
"use strict";
|
||||
|
||||
// ── Compact number formatting ──────────────────────────────────
|
||||
function formatCompact(n) {
|
||||
if (n == null) return null;
|
||||
if (n >= 1_000_000) return (n / 1_000_000).toFixed(1).replace(/\.0$/, "") + "M";
|
||||
if (n >= 1_000) return (n / 1_000).toFixed(1).replace(/\.0$/, "") + "k";
|
||||
return String(n);
|
||||
}
|
||||
|
||||
// ── Badge creation ─────────────────────────────────────────────
|
||||
function createStatBadge(icon, value, title) {
|
||||
if (value == null) return null;
|
||||
const badge = document.createElement("span");
|
||||
badge.className = "lm-stat-badge";
|
||||
badge.title = title;
|
||||
badge.innerHTML = `<i class="fas fa-${icon}"></i> ${formatCompact(value)}`;
|
||||
return badge;
|
||||
}
|
||||
|
||||
// ── Inject CSS ─────────────────────────────────────────────────
|
||||
function injectStyles() {
|
||||
if (document.getElementById("lm-stats-styles")) return;
|
||||
const style = document.createElement("style");
|
||||
style.id = "lm-stats-styles";
|
||||
style.textContent = `
|
||||
.lm-stat-badges {
|
||||
display: flex;
|
||||
gap: 6px;
|
||||
flex-wrap: wrap;
|
||||
padding: 2px 6px;
|
||||
}
|
||||
.lm-stat-badge {
|
||||
display: inline-flex;
|
||||
align-items: center;
|
||||
gap: 3px;
|
||||
font-size: 11px;
|
||||
padding: 1px 5px;
|
||||
border-radius: 4px;
|
||||
background: rgba(255,255,255,0.1);
|
||||
color: rgba(255,255,255,0.8);
|
||||
white-space: nowrap;
|
||||
}
|
||||
.lm-stat-badge i {
|
||||
font-size: 10px;
|
||||
opacity: 0.7;
|
||||
}
|
||||
`;
|
||||
document.head.appendChild(style);
|
||||
}
|
||||
|
||||
// ── Intercept list API to capture stats ────────────────────────
|
||||
const _statsMap = {};
|
||||
|
||||
const _origFetch = window.fetch;
|
||||
window.fetch = async function (...args) {
|
||||
const response = await _origFetch.apply(this, args);
|
||||
const url = typeof args[0] === "string" ? args[0] : args[0]?.url;
|
||||
if (url && url.includes("/api/lm/") && url.includes("/list")) {
|
||||
try {
|
||||
const clone = response.clone();
|
||||
const data = await clone.json();
|
||||
(data.items || []).forEach((item) => {
|
||||
if (item.sha256 && item.download_count != null) {
|
||||
_statsMap[item.sha256] = {
|
||||
download_count: item.download_count,
|
||||
rating: item.rating,
|
||||
rating_count: item.rating_count,
|
||||
thumbs_up_count: item.thumbs_up_count,
|
||||
};
|
||||
}
|
||||
});
|
||||
} catch (e) { /* ignore parse errors */ }
|
||||
}
|
||||
return response;
|
||||
};
|
||||
|
||||
// ── Patch model cards ──────────────────────────────────────────
|
||||
function patchCards() {
|
||||
const cards = document.querySelectorAll(".model-card:not([data-stats-patched])");
|
||||
cards.forEach((card) => {
|
||||
card.setAttribute("data-stats-patched", "1");
|
||||
|
||||
const sha = card.dataset.sha256;
|
||||
if (!sha || !_statsMap[sha]) return;
|
||||
|
||||
const stats = _statsMap[sha];
|
||||
const container = document.createElement("div");
|
||||
container.className = "lm-stat-badges";
|
||||
|
||||
const dlBadge = createStatBadge("download", stats.download_count, "Downloads");
|
||||
const ratingBadge = createStatBadge("star",
|
||||
stats.rating ? Number(stats.rating.toFixed(1)) : null, "Rating");
|
||||
const thumbsBadge = createStatBadge("thumbs-up", stats.thumbs_up_count, "Likes");
|
||||
|
||||
[dlBadge, ratingBadge, thumbsBadge].forEach((b) => {
|
||||
if (b) container.appendChild(b);
|
||||
});
|
||||
|
||||
if (container.children.length > 0) {
|
||||
// Insert inside .model-info, after .model-name
|
||||
const modelInfo = card.querySelector(".model-info");
|
||||
if (modelInfo) {
|
||||
modelInfo.appendChild(container);
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
// ── Patch sort dropdown ────────────────────────────────────────
|
||||
function patchSortDropdown() {
|
||||
const select = document.getElementById("sortSelect");
|
||||
if (!select || select.querySelector('[value="downloads:desc"]')) return;
|
||||
|
||||
const group = document.createElement("optgroup");
|
||||
group.label = "CivitAI Stats";
|
||||
|
||||
const options = [
|
||||
["downloads:desc", "Most downloaded"],
|
||||
["downloads:asc", "Least downloaded"],
|
||||
["rating:desc", "Highest rated"],
|
||||
["rating:asc", "Lowest rated"],
|
||||
["thumbsup:desc", "Most liked"],
|
||||
["thumbsup:asc", "Least liked"],
|
||||
];
|
||||
|
||||
options.forEach(([value, label]) => {
|
||||
const opt = document.createElement("option");
|
||||
opt.value = value;
|
||||
opt.textContent = label;
|
||||
group.appendChild(opt);
|
||||
});
|
||||
|
||||
select.appendChild(group);
|
||||
}
|
||||
|
||||
// ── Toolbar "Fetch Stats" button ───────────────────────────────
|
||||
function addFetchStatsButton() {
|
||||
const toolbar = document.querySelector(".action-buttons");
|
||||
if (!toolbar || document.getElementById("fetchStatsBtn")) return;
|
||||
|
||||
const group = document.createElement("div");
|
||||
group.className = "control-group";
|
||||
group.innerHTML = `
|
||||
<button id="fetchStatsBtn" data-action="fetch-stats"
|
||||
title="Fetch CivitAI stats (downloads, ratings, likes)">
|
||||
<i class="fas fa-chart-bar"></i> <span>Fetch Stats</span>
|
||||
</button>
|
||||
`;
|
||||
|
||||
// Insert before the bulk operations button
|
||||
const bulkBtn = document.getElementById("bulkOperationsBtn");
|
||||
if (bulkBtn && bulkBtn.closest(".control-group")) {
|
||||
toolbar.insertBefore(group, bulkBtn.closest(".control-group"));
|
||||
} else {
|
||||
toolbar.appendChild(group);
|
||||
}
|
||||
|
||||
group.querySelector("button").addEventListener("click", async () => {
|
||||
const btn = document.getElementById("fetchStatsBtn");
|
||||
btn.disabled = true;
|
||||
btn.innerHTML = '<i class="fas fa-spinner fa-spin"></i> <span>Fetching...</span>';
|
||||
|
||||
try {
|
||||
const resp = await _origFetch("/api/lm-extra/fetch-stats", { method: "POST" });
|
||||
const data = await resp.json();
|
||||
if (data.success) {
|
||||
btn.innerHTML = `<i class="fas fa-check"></i> <span>${data.updated} updated</span>`;
|
||||
setTimeout(() => {
|
||||
btn.innerHTML = '<i class="fas fa-chart-bar"></i> <span>Fetch Stats</span>';
|
||||
btn.disabled = false;
|
||||
}, 3000);
|
||||
// Trigger page reload to show new stats
|
||||
const sortSelect = document.getElementById("sortSelect");
|
||||
if (sortSelect) {
|
||||
sortSelect.dispatchEvent(new Event("change"));
|
||||
}
|
||||
} else {
|
||||
throw new Error(data.error || "Unknown error");
|
||||
}
|
||||
} catch (err) {
|
||||
btn.innerHTML = '<i class="fas fa-exclamation-triangle"></i> <span>Error</span>';
|
||||
console.error("[LM-Stats] Fetch failed:", err);
|
||||
setTimeout(() => {
|
||||
btn.innerHTML = '<i class="fas fa-chart-bar"></i> <span>Fetch Stats</span>';
|
||||
btn.disabled = false;
|
||||
}, 3000);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
// ── Observe DOM for card rendering ─────────────────────────────
|
||||
function startObserver() {
|
||||
injectStyles();
|
||||
patchSortDropdown();
|
||||
addFetchStatsButton();
|
||||
patchCards();
|
||||
|
||||
const observer = new MutationObserver(() => {
|
||||
patchCards();
|
||||
patchSortDropdown();
|
||||
addFetchStatsButton();
|
||||
});
|
||||
|
||||
observer.observe(document.body, { childList: true, subtree: true });
|
||||
}
|
||||
|
||||
// ── Init ───────────────────────────────────────────────────────
|
||||
if (document.readyState === "loading") {
|
||||
document.addEventListener("DOMContentLoaded", startObserver);
|
||||
} else {
|
||||
startObserver();
|
||||
}
|
||||
})();
|
||||
Reference in New Issue
Block a user