diff --git a/stats_service.py b/stats_service.py new file mode 100644 index 0000000..e431d25 --- /dev/null +++ b/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 diff --git a/tests/test_stats_service.py b/tests/test_stats_service.py new file mode 100644 index 0000000..963fe9c --- /dev/null +++ b/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