feat: add CivitAI stats fetch service
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -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
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user