feat: add CivitAI stats fetch service

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

stats_service.py

@@ -0,0 +1,142 @@
+"""Service for fetching CivitAI stats and storing them in the local DB."""
+from __future__ import annotations
+
+import asyncio
+import logging
+from typing import Any
+
+import aiohttp
+
+try:
+    from .stats_db import StatsDB
+except ImportError:
+    from stats_db import StatsDB
+
+logger = logging.getLogger(__name__)
+
+_CIVITAI_API = "https://civitai.com/api/v1"
+_RATE_LIMIT_DELAY = 1.5  # seconds between requests
+
+
+def extract_version_stats(model_data: dict) -> list[tuple[str, dict]]:
+    """Extract (sha256, stats_dict) pairs from a CivitAI model response.
+
+    Each model has multiple versions; each version has files with hashes.
+    Returns one entry per version that has a SHA256 hash.
+    """
+    model_id = model_data.get("id")
+    results = []
+
+    for version in model_data.get("modelVersions", []):
+        version_id = version.get("id")
+        sha256 = None
+        for f in version.get("files", []):
+            sha256 = (f.get("hashes") or {}).get("SHA256")
+            if sha256:
+                break
+        if not sha256:
+            continue
+
+        stats = version.get("stats") or {}
+        results.append((
+            sha256.lower(),
+            {
+                "civitai_model_id": model_id,
+                "civitai_version_id": version_id,
+                "download_count": stats.get("downloadCount", 0),
+                "rating": stats.get("rating", 0),
+                "rating_count": stats.get("ratingCount", 0),
+                "thumbs_up_count": stats.get("thumbsUpCount", 0),
+            },
+        ))
+
+    return results
+
+
+class StatsFetchService:
+    """Fetches CivitAI stats for models and stores them in StatsDB."""
+
+    def __init__(self, db: StatsDB, api_key: str | None = None):
+        self.db = db
+        self._api_key = api_key
+        self._session: aiohttp.ClientSession | None = None
+
+    async def _get_session(self) -> aiohttp.ClientSession:
+        if self._session is None or self._session.closed:
+            headers = {}
+            if self._api_key:
+                headers["Authorization"] = f"Bearer {self._api_key}"
+            timeout = aiohttp.ClientTimeout(total=30)
+            self._session = aiohttp.ClientSession(
+                headers=headers, timeout=timeout
+            )
+        return self._session
+
+    async def _fetch_model(self, model_id: int) -> dict | None:
+        """Fetch a single model's data from CivitAI API."""
+        url = f"{_CIVITAI_API}/models/{model_id}"
+        session = await self._get_session()
+        try:
+            async with session.get(url) as resp:
+                if resp.status == 429:
+                    logger.warning("[LM-Remote] CivitAI rate limited, backing off")
+                    await asyncio.sleep(5)
+                    return None
+                if resp.status != 200:
+                    logger.debug("[LM-Remote] CivitAI returned %d for model %d",
+                                 resp.status, model_id)
+                    return None
+                return await resp.json()
+        except Exception as exc:
+            logger.warning("[LM-Remote] CivitAI fetch failed for model %d: %s",
+                           model_id, exc)
+            return None
+
+    async def fetch_stats_for_models(
+        self,
+        models: list[dict],
+        progress_callback: Any = None,
+    ) -> int:
+        """Fetch stats from CivitAI for a list of models.
+
+        Args:
+            models: list of dicts with 'sha256' and 'civitai_model_id' keys.
+            progress_callback: optional async callable(current, total).
+
+        Returns:
+            Number of model versions successfully updated.
+        """
+        # Deduplicate by model_id (multiple versions share the same model)
+        seen_model_ids: set[int] = set()
+        unique_models: list[dict] = []
+        for m in models:
+            mid = m.get("civitai_model_id")
+            if mid and mid not in seen_model_ids:
+                seen_model_ids.add(mid)
+                unique_models.append(m)
+
+        total = len(unique_models)
+        updated = 0
+
+        for i, model in enumerate(unique_models):
+            model_id = model["civitai_model_id"]
+            data = await self._fetch_model(model_id)
+            if data:
+                rows = extract_version_stats(data)
+                if rows:
+                    self.db.upsert_batch(rows)
+                    updated += len(rows)
+
+            if progress_callback:
+                await progress_callback(i + 1, total)
+
+            # Rate limiting between requests
+            if i < total - 1:
+                await asyncio.sleep(_RATE_LIMIT_DELAY)
+
+        return updated
+
+    async def close(self) -> None:
+        if self._session and not self._session.closed:
+            await self._session.close()
+            self._session = None

tests/test_stats_service.py

@@ -0,0 +1,121 @@
+"""Tests for the CivitAI stats fetch service."""
+import pytest
+from unittest.mock import patch
+
+
+@pytest.fixture
+def db(tmp_path):
+    from stats_db import StatsDB
+    db = StatsDB(tmp_path / "test.db")
+    db.init()
+    return db
+
+
+@pytest.fixture
+def service(db):
+    from stats_service import StatsFetchService
+    return StatsFetchService(db)
+
+
+def test_extract_version_stats():
+    """Extract stats from a CivitAI model API response."""
+    from stats_service import extract_version_stats
+
+    model_response = {
+        "id": 123,
+        "modelVersions": [
+            {
+                "id": 456,
+                "files": [{"hashes": {"SHA256": "AABB"}}],
+                "stats": {
+                    "downloadCount": 5000,
+                    "ratingCount": 120,
+                    "rating": 4.8,
+                    "thumbsUpCount": 89,
+                },
+            },
+            {
+                "id": 789,
+                "files": [{"hashes": {"SHA256": "CCDD"}}],
+                "stats": {
+                    "downloadCount": 1000,
+                    "ratingCount": 30,
+                    "rating": 3.5,
+                    "thumbsUpCount": 20,
+                },
+            },
+        ],
+    }
+    result = extract_version_stats(model_response)
+    assert len(result) == 2
+    assert result[0][0] == "aabb"  # sha256 lowercased
+    assert result[0][1]["download_count"] == 5000
+    assert result[0][1]["civitai_model_id"] == 123
+    assert result[0][1]["civitai_version_id"] == 456
+    assert result[1][0] == "ccdd"
+
+
+def test_extract_version_stats_no_hash():
+    """Versions without a SHA256 hash are skipped."""
+    from stats_service import extract_version_stats
+
+    model_response = {
+        "id": 1,
+        "modelVersions": [
+            {"id": 10, "files": [{"hashes": {}}], "stats": {"downloadCount": 1}},
+        ],
+    }
+    assert extract_version_stats(model_response) == []
+
+
+def test_extract_version_stats_no_stats():
+    """Versions without stats get zeros."""
+    from stats_service import extract_version_stats
+
+    model_response = {
+        "id": 1,
+        "modelVersions": [
+            {"id": 10, "files": [{"hashes": {"SHA256": "AABB"}}]},
+        ],
+    }
+    result = extract_version_stats(model_response)
+    assert len(result) == 1
+    assert result[0][1]["download_count"] == 0
+
+
+@pytest.mark.asyncio
+async def test_fetch_stats_for_models(service):
+    """fetch_stats_for_models calls CivitAI API and stores results."""
+    models = [
+        {"sha256": "aabb", "civitai_model_id": 123},
+        {"sha256": "ccdd", "civitai_model_id": 456},
+    ]
+    mock_response_123 = {
+        "id": 123,
+        "modelVersions": [{
+            "id": 10,
+            "files": [{"hashes": {"SHA256": "AABB"}}],
+            "stats": {"downloadCount": 500, "ratingCount": 10,
+                      "rating": 4.0, "thumbsUpCount": 8},
+        }],
+    }
+    mock_response_456 = {
+        "id": 456,
+        "modelVersions": [{
+            "id": 20,
+            "files": [{"hashes": {"SHA256": "CCDD"}}],
+            "stats": {"downloadCount": 100, "ratingCount": 5,
+                      "rating": 3.0, "thumbsUpCount": 2},
+        }],
+    }
+
+    async def mock_fetch(model_id):
+        return {123: mock_response_123, 456: mock_response_456}.get(model_id)
+
+    with patch.object(service, "_fetch_model", side_effect=mock_fetch):
+        count = await service.fetch_stats_for_models(models)
+
+    assert count == 2
+    result = service.db.get_by_hashes(["aabb", "ccdd"])
+    assert result["aabb"]["download_count"] == 500
+    assert result["ccdd"]["download_count"] == 100