feat: enrich list API responses with CivitAI stats from local DB

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

proxy.py

@@ -11,13 +11,21 @@ to the local ComfyUI frontend (the remote instance has no connected browsers).
 from __future__ import annotations
 
 import asyncio
+import json
 import logging
 
 import aiohttp
 from aiohttp import web, WSMsgType
 
+try:
     from .config import remote_config
+except ImportError:
+    from config import remote_config  # type: ignore[no-redef]
+
+try:
     from .remote_client import RemoteLoraClient
+except ImportError:
+    from remote_client import RemoteLoraClient  # type: ignore[no-redef]
 
 logger = logging.getLogger(__name__)
 
@@ -178,6 +186,53 @@ _SEND_SYNC_HANDLERS = {
     "/api/lm/register-nodes": _handle_register_nodes,
 }
 
+# ---------------------------------------------------------------------------
+# Stats enrichment
+# ---------------------------------------------------------------------------
+
+_stats_db = None
+
+
+def _get_stats_db():
+    """Lazy-initialize the stats database."""
+    global _stats_db
+    if _stats_db is None:
+        try:
+            from .stats_db import StatsDB as _StatsDBClass
+        except ImportError:
+            from stats_db import StatsDB as _StatsDBClass
+        _stats_db = _StatsDBClass()
+        _stats_db.init()
+    return _stats_db
+
+
+def _enrich_list_response(response_data: dict, db=None) -> dict:
+    """Merge stats from the local DB into a list API response."""
+    items = response_data.get("items", [])
+    if not items:
+        return response_data
+
+    if db is None:
+        db = _get_stats_db()
+
+    hashes = [item["sha256"] for item in items if item.get("sha256")]
+    if not hashes:
+        return response_data
+
+    stats_map = db.get_by_hashes(hashes)
+
+    for item in items:
+        sha = item.get("sha256")
+        if sha and sha in stats_map:
+            s = stats_map[sha]
+            item["download_count"] = s["download_count"]
+            item["rating"] = s["rating"]
+            item["rating_count"] = s["rating_count"]
+            item["thumbs_up_count"] = s["thumbs_up_count"]
+
+    return response_data
+
+
 # Shared HTTP session for proxied requests (connection pooling)
 _proxy_session: aiohttp.ClientSession | None = None
 
@@ -332,7 +387,17 @@ async def lm_remote_proxy_middleware(request: web.Request, handler):
 
     # Regular proxy routes
     if _should_proxy(path):
-        return await _proxy_http(request)
+        response = await _proxy_http(request)
+        # Enrich model list responses with stats from local DB
+        if (path.endswith("/list") and path.startswith("/api/lm/")
+                and response.status == 200):
+            try:
+                body = json.loads(response.body)
+                enriched = _enrich_list_response(body)
+                return web.json_response(enriched)
+            except Exception:
+                pass  # Return original response on any error
+        return response
 
     # Not a LoRA Manager route — fall through
     return await handler(request)

remote_client.py

@@ -8,7 +8,10 @@ from typing import Any
 
 import aiohttp
 
+try:
     from .config import remote_config
+except ImportError:
+    from config import remote_config  # type: ignore[no-redef]
 
 logger = logging.getLogger(__name__)
 

tests/test_proxy_enrichment.py

@@ -0,0 +1,41 @@
+"""Tests for proxy response enrichment with stats."""
+import pytest
+
+
+@pytest.fixture
+def db(tmp_path):
+    from stats_db import StatsDB
+    db = StatsDB(tmp_path / "test.db")
+    db.init()
+    return db
+
+
+def test_enrich_list_response(db):
+    """Stats are merged into list response items."""
+    from proxy import _enrich_list_response
+
+    db.upsert("hash_a", {"download_count": 5000, "rating": 4.8,
+                          "rating_count": 120, "thumbs_up_count": 89})
+
+    response_data = {
+        "items": [
+            {"sha256": "hash_a", "file_name": "lora_a.safetensors"},
+            {"sha256": "hash_b", "file_name": "lora_b.safetensors"},
+        ],
+        "total": 2,
+    }
+
+    enriched = _enrich_list_response(response_data, db)
+    assert enriched["items"][0]["download_count"] == 5000
+    assert enriched["items"][0]["rating"] == 4.8
+    assert enriched["items"][0]["thumbs_up_count"] == 89
+    # Item without stats gets no extra fields
+    assert "download_count" not in enriched["items"][1]
+
+
+def test_enrich_empty_items(db):
+    """Empty items list returns unchanged."""
+    from proxy import _enrich_list_response
+    response_data = {"items": [], "total": 0}
+    enriched = _enrich_list_response(response_data, db)
+    assert enriched["items"] == []