diff --git a/proxy.py b/proxy.py index 78bc180..f95b773 100644 --- a/proxy.py +++ b/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 -from .config import remote_config -from .remote_client import RemoteLoraClient +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) diff --git a/remote_client.py b/remote_client.py index 35f4d83..ecfa4ee 100644 --- a/remote_client.py +++ b/remote_client.py @@ -8,7 +8,10 @@ from typing import Any import aiohttp -from .config import remote_config +try: + from .config import remote_config +except ImportError: + from config import remote_config # type: ignore[no-redef] logger = logging.getLogger(__name__) diff --git a/tests/test_proxy_enrichment.py b/tests/test_proxy_enrichment.py new file mode 100644 index 0000000..88459d2 --- /dev/null +++ b/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"] == []