alphard/ComfyUI-Lora-Manager
1
0
Fork 0
mirror of https://github.com/willmiao/ComfyUI-Lora-Manager.git synced 2026-10-05 09:35:31 -03:00
Code Issues Packages Projects Releases Wiki Activity
Files
ad2402724b2edce9ac720644a5022ed225865967
ComfyUI-Lora-Manager/py/services/model_update_service.py
T
Will Miao ad2402724b feat(update): track buzz prices and alert below a threshold 
CivitAI's public API deliberately omits prices — paidAccess is trimmed to
{permanent, endsAt} because "pricing belongs to the purchase flow" — but the
public model page embeds the site's own model.getById result, including
paidAccess.terms, in its server-rendered payload. That is read anonymously
(no API key, no internal endpoint, no forged Origin), one request per gated
model, so only the ~2% of models that actually carry a gate pay for it.

- optional capture, off by default: price_tracking_enabled,
  price_alert_threshold_buzz (0 = alert on "became free" only) and
  price_check_ttl_hours; prices refresh on their own TTL and immediately when a
  gate changes, and a failed fetch keeps the stored price instead of blanking it
- versions that stop carrying a gate are marked free (persisted gate_lapsed_at)
  and gate transitions are reported as events on the refresh response, so a
  version already in the library can announce that it became free
- price_alert_state plus a price_drop edge event; new
  GET /api/lm/{type}/updates/price-alerts lists what is under the threshold
- versions tab shows the price (effective, with the list price struck through
  and a Blue Buzz note) and a Free Now badge; an update check toasts the
  transitions in one message
- the parser and the alerts query are unit-tested against a trimmed page
  fixture, and every route definition is now asserted to resolve to a handler

Plan, verification notes and the deviations from it are in
docs/plans/paid-model-price-tracking.md.
2026-10-04 08:53:24 +08:00

2516 lines
96 KiB
Python
Raw Blame History

# pyright: reportImportCycles=false
# Lazy (function-local) imports still count as static edges in basedpyright's
# reportImportCycles, so the ServiceRegistry singleton pattern necessarily forms
# import cycles. Breaking them would require an architectural refactor.
"""Service for tracking remote model version updates."""
from __future__ import annotations
import asyncio
import json
import logging
import os
import sqlite3
import time
from dataclasses import dataclass, field, replace
from datetime import datetime, timezone
from typing import Any, Dict, Iterable, Iterator, List, Mapping, Optional, Sequence
from .errors import RateLimitError, ResourceNotFoundError
from .settings_manager import get_settings_manager
from ..utils.cache_paths import CacheType, resolve_cache_path_with_migration
from ..utils.constants import MODEL_WEIGHT_FILE_TYPES
from ..utils.civitai_utils import rewrite_preview_url
from ..utils.paid_access import (
is_early_access_deadline_active,
normalize_paid_access as _normalize_paid_access_payload,
)
from ..utils.preview_selection import resolve_mature_threshold, select_preview_media
logger = logging.getLogger(__name__)
# Version fields a price payload may set (see py/utils/civitai_page_prices.py).
_PRICE_FIELD_NAMES = frozenset(
{
"price_buzz",
"list_price_buzz",
"generation_price_buzz",
"accepts_blue_buzz",
"price_sale_ends_at",
}
)
def _normalize_int(value) -> Optional[int]:
"""Safely convert a value to an integer."""
try:
if value is None:
return None
return int(value)
except (TypeError, ValueError):
return None
def _format_utc_timestamp(value: float) -> str:
"""Format a POSIX timestamp the way CivitAI serializes its timestamps."""
return (
datetime.fromtimestamp(value, tz=timezone.utc)
.isoformat()
.replace("+00:00", "Z")
)
def _normalize_string(value) -> Optional[str]:
"""Return a stripped string or None if the value is empty."""
if value is None:
return None
if isinstance(value, str):
stripped = value.strip()
return stripped or None
try:
normalized = str(value).strip()
return normalized or None
except Exception:
return None
def _normalize_base_model(value) -> Optional[str]:
"""Normalize base-model names for case-insensitive comparison."""
normalized = _normalize_string(value)
if normalized is None:
return None
return normalized.lower()
@dataclass
class ModelVersionRecord:
"""Persisted metadata for a single model version."""
version_id: int
name: Optional[str]
base_model: Optional[str]
released_at: Optional[str]
size_bytes: Optional[int]
preview_url: Optional[str]
is_in_library: bool
should_ignore: bool
early_access_ends_at: Optional[str] = None
sort_index: int = 0
is_early_access: bool = False
usage_control: Optional[str] = None # "Download", "Generation", "InternalGeneration"
paid_access: Optional[str] = None # JSON string of the CivitAI paidAccess DTO
is_paid: bool = False # True when paidAccess.permanent is True (permanent paid gate)
# Number of downloadable weight files for the version (None when unknown,
# e.g. records persisted before this field existed or locally-synthesized
# entries). Mirrors the frontend isModelWeightFile() filter.
file_count: Optional[int] = None
# When a version that used to carry a paid gate stopped carrying one. Kept so
# the UI can mark a version "free" long after the transition happened; cleared
# again if the version is gated once more.
gate_lapsed_at: Optional[str] = None
# Download price in Buzz for the gated version, and friends. All None until a
# price refresh succeeds (see the price tracking feature); `price_buzz` is the
# effective price a buyer pays now (sale-adjusted) and `list_price_buzz` the
# undiscounted stored price.
price_buzz: Optional[int] = None
list_price_buzz: Optional[int] = None
generation_price_buzz: Optional[int] = None
accepts_blue_buzz: bool = False
price_sale_ends_at: Optional[str] = None
price_checked_at: Optional[float] = None
# Last computed "price alert threshold hit" state, used to detect the edge
# (0 -> 1) that fires a one-shot price_drop event.
price_alert_state: bool = False
@dataclass
class ModelUpdateRecord:
"""Representation of a persisted update record."""
model_type: str
model_id: int
versions: List[ModelVersionRecord]
last_checked_at: Optional[float]
should_ignore_model: bool
# Gate-state and price transitions observed during the refresh that produced
# this record. Derived, never persisted: a record read back from SQLite has an
# empty list. Each entry is {"versionId", "kind", "versionName",
# "isInLibrary", "priceBuzz"?} with kind in
# {"new_gate", "became_free", "price_drop"}.
events: List[Dict[str, Any]] = field(default_factory=list)
@property
def largest_version_id(self) -> Optional[int]:
"""Return the highest known version identifier for the model."""
if not self.versions:
return None
return max(version.version_id for version in self.versions)
@property
def version_ids(self) -> List[int]:
"""Return all known version identifiers."""
return [version.version_id for version in self.versions]
@property
def in_library_version_ids(self) -> List[int]:
"""Return the subset of version identifiers present in the local library."""
return [version.version_id for version in self.versions if version.is_in_library]
def has_update(
self,
hide_early_access: bool = False,
hide_non_downloadable: bool = True,
hide_paid: bool = False,
) -> bool:
"""Return True when a non-ignored remote version newer than the newest local copy is available.
Args:
hide_early_access: If True, exclude early access versions from update check.
hide_non_downloadable: If True, exclude versions that don't allow downloads.
hide_paid: If True, exclude permanent paid versions from update check.
"""
if self.should_ignore_model:
return False
max_in_library = None
for version in self.versions:
if version.is_in_library:
if max_in_library is None or version.version_id > max_in_library:
max_in_library = version.version_id
if max_in_library is None:
return any(
not version.is_in_library
and not version.should_ignore
and not (hide_early_access and ModelUpdateRecord._is_early_access_active(version))
and not (hide_paid and version.is_paid)
and not (hide_non_downloadable and not ModelUpdateRecord._is_downloadable(version))
for version in self.versions
)
for version in self.versions:
if version.is_in_library or version.should_ignore:
continue
if hide_early_access and ModelUpdateRecord._is_early_access_active(version):
continue
if hide_paid and version.is_paid:
continue
if hide_non_downloadable and not ModelUpdateRecord._is_downloadable(version):
continue
if version.version_id > max_in_library:
return True
return False
@staticmethod
def _is_early_access_active(version: ModelVersionRecord) -> bool:
"""Check if a version is currently in early access period.
Uses two-phase detection:
1. If exact EA end time available (from single version API), use it for precise check
2. Otherwise fallback to basic EA flag (from bulk API)
"""
# Permanent paid versions are not early access; they are filtered by
# hide_paid instead. Only timed gates count as early access.
if version.is_paid and not version.early_access_ends_at:
return False
# Phase 2: Precise check with exact end time (None when the gate is timed
# but its window end has not been recorded yet -> treated as active below)
if version.early_access_ends_at:
return is_early_access_deadline_active(version.early_access_ends_at)
# Phase 1: Basic EA flag from bulk API
return version.is_early_access
@staticmethod
def _is_downloadable(version: ModelVersionRecord) -> bool:
if version.usage_control is None:
return True
return version.usage_control == "Download"
def has_update_for_base(
self,
local_version_id: Optional[int],
local_base_model: Optional[str],
hide_early_access: bool = False,
hide_non_downloadable: bool = True,
hide_paid: bool = False,
) -> bool:
"""Return True when a newer remote version with the same base model exists.
Args:
local_version_id: The current local version id.
local_base_model: The base model to filter by.
hide_early_access: If True, exclude early access versions from update check.
hide_non_downloadable: If True, exclude versions that don't allow downloads.
hide_paid: If True, exclude permanent paid versions from update check.
"""
if self.should_ignore_model:
return False
normalized_base = _normalize_base_model(local_base_model)
if normalized_base is None:
return False
threshold = _normalize_int(local_version_id)
if threshold is None:
highest_local = None
for version in self.versions:
if not version.is_in_library:
continue
version_base = _normalize_base_model(version.base_model)
if version_base != normalized_base:
continue
if highest_local is None or version.version_id > highest_local:
highest_local = version.version_id
threshold = highest_local
if threshold is None:
return False
for version in self.versions:
if version.is_in_library or version.should_ignore:
continue
if hide_early_access and ModelUpdateRecord._is_early_access_active(version):
continue
if hide_paid and version.is_paid:
continue
if hide_non_downloadable and not ModelUpdateRecord._is_downloadable(version):
continue
version_base = _normalize_base_model(version.base_model)
if version_base != normalized_base:
continue
if version.version_id > threshold:
return True
return False
def has_update_for_local_bases(
self,
hide_early_access: bool = False,
hide_non_downloadable: bool = True,
hide_paid: bool = False,
) -> bool:
"""Return True when any locally-held base model scope has an update.
Aggregates :meth:`has_update_for_base` across every distinct base model
present among in-library versions. This mirrors the per-item evaluation
performed by ``BaseModelService._annotate_update_flags`` when the
``version_grouping`` setting is ``same_base``, so callers reporting
"how many models have updates" stay aligned with what the Updates
filter displays. Use this instead of :meth:`has_update` for such
summaries; see issue #1083.
When no local base model is known (nothing held locally, or versions
never seen in any remote listing), falls back to :meth:`has_update` so
a model the item-level filter may still flag is not silently dropped
from summaries.
"""
bases = {
_normalize_base_model(version.base_model)
for version in self.versions
if version.is_in_library
}
bases.discard(None)
if not bases:
return self.has_update(
hide_early_access=hide_early_access,
hide_non_downloadable=hide_non_downloadable,
hide_paid=hide_paid,
)
return any(
self.has_update_for_base(
None,
base,
hide_early_access=hide_early_access,
hide_non_downloadable=hide_non_downloadable,
hide_paid=hide_paid,
)
for base in bases
)
class ModelUpdateService:
"""Persist and query remote model version metadata."""
_SQLITE_MAX_VARIABLES = 500
_SCHEMA = """
PRAGMA foreign_keys = ON;
CREATE TABLE IF NOT EXISTS model_update_status (
model_id INTEGER PRIMARY KEY,
model_type TEXT NOT NULL,
last_checked_at REAL,
should_ignore_model INTEGER NOT NULL DEFAULT 0
);
CREATE TABLE IF NOT EXISTS model_update_versions (
model_id INTEGER NOT NULL,
version_id INTEGER NOT NULL,
sort_index INTEGER NOT NULL DEFAULT 0,
name TEXT,
base_model TEXT,
released_at TEXT,
size_bytes INTEGER,
preview_url TEXT,
is_in_library INTEGER NOT NULL DEFAULT 0,
should_ignore INTEGER NOT NULL DEFAULT 0,
usage_control TEXT,
paid_access TEXT,
is_paid INTEGER NOT NULL DEFAULT 0,
file_count INTEGER,
gate_lapsed_at TEXT,
price_buzz INTEGER,
list_price_buzz INTEGER,
generation_price_buzz INTEGER,
accepts_blue_buzz INTEGER NOT NULL DEFAULT 0,
price_sale_ends_at TEXT,
price_checked_at REAL,
price_alert_state INTEGER NOT NULL DEFAULT 0,
PRIMARY KEY (model_id, version_id),
FOREIGN KEY(model_id) REFERENCES model_update_status(model_id) ON DELETE CASCADE
);
CREATE INDEX IF NOT EXISTS idx_model_update_versions_model_id
ON model_update_versions(model_id);
"""
def __init__(
self,
db_path: str | None = None,
*,
ttl_seconds: int = 24 * 60 * 60,
settings_manager=None,
) -> None:
self._settings = settings_manager or get_settings_manager()
self._library_name = self._get_active_library_name()
self._db_path = db_path or self._resolve_default_path(self._library_name)
self._ttl_seconds = ttl_seconds
self._lock = asyncio.Lock()
self._schema_initialized = False
self._custom_db_path = db_path is not None
self._ensure_directory()
self._initialize_schema()
def _get_active_library_name(self) -> str:
try:
value = self._settings.get_active_library_name()
except Exception:
value = None
return value or "default"
def _resolve_default_path(self, library_name: str) -> str:
env_override = os.environ.get("LORA_MANAGER_MODEL_UPDATE_DB")
return resolve_cache_path_with_migration(
CacheType.MODEL_UPDATE,
library_name=library_name,
env_override=env_override,
)
def on_library_changed(self) -> None:
"""Switch to the database for the active library."""
if self._custom_db_path:
return
library_name = self._get_active_library_name()
new_path = self._resolve_default_path(library_name)
if new_path == self._db_path:
return
self._library_name = library_name
self._db_path = new_path
self._schema_initialized = False
self._ensure_directory()
self._initialize_schema()
def _ensure_directory(self) -> None:
directory = os.path.dirname(self._db_path)
if directory:
os.makedirs(directory, exist_ok=True)
def _connect(self) -> sqlite3.Connection:
conn = sqlite3.connect(self._db_path, check_same_thread=False)
conn.row_factory = sqlite3.Row
return conn
def _initialize_schema(self) -> None:
if self._schema_initialized:
return
try:
with self._connect() as conn:
conn.execute("PRAGMA journal_mode=WAL")
conn.execute("PRAGMA foreign_keys = ON")
conn.executescript(self._SCHEMA)
self._apply_migrations(conn)
self._migrate_from_legacy_snapshot(conn)
self._schema_initialized = True
except Exception as exc: # pragma: no cover - defensive guard
logger.error("Failed to initialize update schema: %s", exc, exc_info=True)
raise
def _migrate_from_legacy_snapshot(self, conn: sqlite3.Connection) -> None:
"""Copy update tracking data out of the legacy model snapshot database."""
if self._custom_db_path:
return
try:
from .persistent_model_cache import PersistentModelCache
legacy_path = PersistentModelCache.get_default(self._library_name).get_database_path()
except Exception:
return
if not legacy_path or os.path.abspath(legacy_path) == os.path.abspath(self._db_path):
return
if not os.path.exists(legacy_path):
return
try:
existing_row = conn.execute(
"SELECT 1 FROM model_update_status LIMIT 1"
).fetchone()
if existing_row:
return
except Exception:
return
try:
with sqlite3.connect(legacy_path, check_same_thread=False) as legacy_conn:
legacy_conn.row_factory = sqlite3.Row
status_rows = legacy_conn.execute(
"""
SELECT model_id, model_type, last_checked_at, should_ignore_model
FROM model_update_status
"""
).fetchall()
if not status_rows:
return
version_rows = legacy_conn.execute(
"""
SELECT model_id, version_id, sort_index, name, base_model, released_at,
size_bytes, preview_url, is_in_library, should_ignore,
early_access_ends_at, is_early_access
FROM model_update_versions
ORDER BY model_id ASC, sort_index ASC, version_id ASC
"""
).fetchall()
conn.execute("BEGIN")
conn.executemany(
"""
INSERT OR REPLACE INTO model_update_status (
model_id, model_type, last_checked_at, should_ignore_model
) VALUES (?, ?, ?, ?)
""",
[
(
int(row["model_id"]),
row["model_type"],
row["last_checked_at"],
int(row["should_ignore_model"] or 0),
)
for row in status_rows
],
)
conn.executemany(
"""
INSERT OR REPLACE INTO model_update_versions (
model_id, version_id, sort_index, name, base_model, released_at,
size_bytes, preview_url, is_in_library, should_ignore,
early_access_ends_at, is_early_access
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
""",
[
(
int(row["model_id"]),
int(row["version_id"]),
int(row["sort_index"] or 0),
row["name"],
row["base_model"],
row["released_at"],
row["size_bytes"],
row["preview_url"],
int(row["is_in_library"] or 0),
int(row["should_ignore"] or 0),
row["early_access_ends_at"],
int(row["is_early_access"] or 0),
)
for row in version_rows
],
)
conn.commit()
logger.info(
"Migrated model update tracking data from legacy snapshot DB for %s",
self._library_name,
)
except sqlite3.OperationalError as exc:
logger.debug("Legacy model update migration skipped: %s", exc)
except Exception as exc: # pragma: no cover - defensive guard
logger.warning("Failed to migrate model update data: %s", exc, exc_info=True)
def _apply_migrations(self, conn: sqlite3.Connection) -> None:
"""Ensure legacy databases match the current schema without dropping data."""
status_columns = self._get_table_columns(conn, "model_update_status")
if "should_ignore_model" not in status_columns:
conn.execute(
"ALTER TABLE model_update_status "
"ADD COLUMN should_ignore_model INTEGER NOT NULL DEFAULT 0"
)
version_columns = self._get_table_columns(conn, "model_update_versions")
migrations = {
"sort_index": (
"ALTER TABLE model_update_versions "
"ADD COLUMN sort_index INTEGER NOT NULL DEFAULT 0"
),
"name": (
"ALTER TABLE model_update_versions "
"ADD COLUMN name TEXT"
),
"base_model": (
"ALTER TABLE model_update_versions "
"ADD COLUMN base_model TEXT"
),
"released_at": (
"ALTER TABLE model_update_versions "
"ADD COLUMN released_at TEXT"
),
"size_bytes": (
"ALTER TABLE model_update_versions "
"ADD COLUMN size_bytes INTEGER"
),
"preview_url": (
"ALTER TABLE model_update_versions "
"ADD COLUMN preview_url TEXT"
),
"is_in_library": (
"ALTER TABLE model_update_versions "
"ADD COLUMN is_in_library INTEGER NOT NULL DEFAULT 0"
),
"should_ignore": (
"ALTER TABLE model_update_versions "
"ADD COLUMN should_ignore INTEGER NOT NULL DEFAULT 0"
),
"early_access_ends_at": (
"ALTER TABLE model_update_versions "
"ADD COLUMN early_access_ends_at TEXT"
),
"is_early_access": (
"ALTER TABLE model_update_versions "
"ADD COLUMN is_early_access INTEGER NOT NULL DEFAULT 0"
),
"usage_control": (
"ALTER TABLE model_update_versions "
"ADD COLUMN usage_control TEXT"
),
"paid_access": (
"ALTER TABLE model_update_versions "
"ADD COLUMN paid_access TEXT"
),
"is_paid": (
"ALTER TABLE model_update_versions "
"ADD COLUMN is_paid INTEGER NOT NULL DEFAULT 0"
),
"file_count": (
"ALTER TABLE model_update_versions "
"ADD COLUMN file_count INTEGER"
),
"gate_lapsed_at": (
"ALTER TABLE model_update_versions "
"ADD COLUMN gate_lapsed_at TEXT"
),
"price_buzz": (
"ALTER TABLE model_update_versions "
"ADD COLUMN price_buzz INTEGER"
),
"list_price_buzz": (
"ALTER TABLE model_update_versions "
"ADD COLUMN list_price_buzz INTEGER"
),
"generation_price_buzz": (
"ALTER TABLE model_update_versions "
"ADD COLUMN generation_price_buzz INTEGER"
),
"accepts_blue_buzz": (
"ALTER TABLE model_update_versions "
"ADD COLUMN accepts_blue_buzz INTEGER NOT NULL DEFAULT 0"
),
"price_sale_ends_at": (
"ALTER TABLE model_update_versions "
"ADD COLUMN price_sale_ends_at TEXT"
),
"price_checked_at": (
"ALTER TABLE model_update_versions "
"ADD COLUMN price_checked_at REAL"
),
"price_alert_state": (
"ALTER TABLE model_update_versions "
"ADD COLUMN price_alert_state INTEGER NOT NULL DEFAULT 0"
),
}
for column, statement in migrations.items():
if column not in version_columns:
conn.execute(statement)
# Refresh column metadata after applying additive migrations.
version_columns = self._get_table_columns(conn, "model_update_versions")
if self._requires_model_update_versions_pk_migration(conn):
self._migrate_model_update_versions_primary_key(
conn, version_columns
)
version_columns = self._get_table_columns(conn, "model_update_versions")
if not self._has_unique_constraint(conn, "model_update_status", "model_id"):
self._deduplicate_model_update_status(conn)
conn.execute(
"CREATE UNIQUE INDEX IF NOT EXISTS "
"uq_model_update_status_model_id ON model_update_status(model_id)"
)
conn.execute(
"CREATE INDEX IF NOT EXISTS idx_model_update_versions_model_id "
"ON model_update_versions(model_id)"
)
def _get_table_columns(self, conn: sqlite3.Connection, table: str) -> set[str]:
"""Return the set of existing columns for a table."""
cursor = conn.execute(f"PRAGMA table_info({table})")
return {row["name"] for row in cursor.fetchall()}
def _has_unique_constraint(
self, conn: sqlite3.Connection, table: str, column: str
) -> bool:
"""Return True when the column already enforces uniqueness."""
cursor = conn.execute(f"PRAGMA table_info({table})")
rows = cursor.fetchall()
column_info = next((row for row in rows if row["name"] == column), None)
if column_info is None:
return False
if column_info["pk"] == 1 and all(
other["pk"] == 0 for other in rows if other["name"] != column
):
return True
index_list = conn.execute(f"PRAGMA index_list({table})").fetchall()
for index in index_list:
if not index["unique"]:
continue
index_name = index["name"]
index_info = conn.execute(f"PRAGMA index_info({index_name})").fetchall()
if len(index_info) == 1 and index_info[0]["name"] == column:
return True
return False
def _requires_model_update_versions_pk_migration(
self, conn: sqlite3.Connection
) -> bool:
"""Detect legacy schemas where version_id is the sole primary key."""
info = conn.execute("PRAGMA table_info(model_update_versions)").fetchall()
pk_columns = [row for row in info if row["pk"]]
if not pk_columns:
return True
if len(pk_columns) == 1:
return pk_columns[0]["name"] == "version_id"
ordered = sorted(pk_columns, key=lambda row: row["pk"])
expected = ["model_id", "version_id"]
return [row["name"] for row in ordered] != expected
def _migrate_model_update_versions_primary_key(
self, conn: sqlite3.Connection, legacy_columns: set[str]
) -> None:
"""Upgrade the versions table to use a composite primary key."""
logger.info("Migrating model_update_versions table to composite primary key")
conn.execute(
"ALTER TABLE model_update_versions RENAME TO model_update_versions_legacy"
)
conn.execute(
"""
CREATE TABLE model_update_versions_new (
model_id INTEGER NOT NULL,
version_id INTEGER NOT NULL,
sort_index INTEGER NOT NULL DEFAULT 0,
name TEXT,
base_model TEXT,
released_at TEXT,
size_bytes INTEGER,
preview_url TEXT,
is_in_library INTEGER NOT NULL DEFAULT 0,
should_ignore INTEGER NOT NULL DEFAULT 0,
early_access_ends_at TEXT,
is_early_access INTEGER NOT NULL DEFAULT 0,
paid_access TEXT,
is_paid INTEGER NOT NULL DEFAULT 0,
file_count INTEGER,
gate_lapsed_at TEXT,
price_buzz INTEGER,
list_price_buzz INTEGER,
generation_price_buzz INTEGER,
accepts_blue_buzz INTEGER NOT NULL DEFAULT 0,
price_sale_ends_at TEXT,
price_checked_at REAL,
price_alert_state INTEGER NOT NULL DEFAULT 0,
PRIMARY KEY (model_id, version_id),
FOREIGN KEY(model_id) REFERENCES model_update_status(model_id) ON DELETE CASCADE
)
"""
)
target_columns = [
"model_id",
"version_id",
"sort_index",
"name",
"base_model",
"released_at",
"size_bytes",
"preview_url",
"is_in_library",
"should_ignore",
"early_access_ends_at",
"is_early_access",
"paid_access",
"is_paid",
"file_count",
"gate_lapsed_at",
"price_buzz",
"list_price_buzz",
"generation_price_buzz",
"accepts_blue_buzz",
"price_sale_ends_at",
"price_checked_at",
"price_alert_state",
]
defaults = {
"sort_index": "0",
"name": "NULL",
"base_model": "NULL",
"released_at": "NULL",
"size_bytes": "NULL",
"preview_url": "NULL",
"is_in_library": "0",
"should_ignore": "0",
"early_access_ends_at": "NULL",
"is_early_access": "0",
"paid_access": "NULL",
"is_paid": "0",
"file_count": "NULL",
"gate_lapsed_at": "NULL",
"price_buzz": "NULL",
"list_price_buzz": "NULL",
"generation_price_buzz": "NULL",
"accepts_blue_buzz": "0",
"price_sale_ends_at": "NULL",
"price_checked_at": "NULL",
"price_alert_state": "0",
}
select_parts = []
for column in target_columns:
if column in legacy_columns:
if column in {"sort_index", "is_in_library", "should_ignore"}:
select_parts.append(f"COALESCE({column}, {defaults[column]})")
else:
select_parts.append(column)
else:
select_parts.append(defaults.get(column, "NULL"))
conn.execute(
"""
INSERT INTO model_update_versions_new ({columns})
SELECT {select_clause}
FROM model_update_versions_legacy
""".format(
columns=", ".join(target_columns),
select_clause=", ".join(select_parts),
)
)
conn.execute("DROP TABLE model_update_versions_legacy")
conn.execute(
"ALTER TABLE model_update_versions_new RENAME TO model_update_versions"
)
def _deduplicate_model_update_status(self, conn: sqlite3.Connection) -> None:
"""Remove duplicate status rows before applying uniqueness constraints."""
duplicates = conn.execute(
"""
SELECT model_id
FROM model_update_status
GROUP BY model_id
HAVING COUNT(*) > 1
"""
).fetchall()
if not duplicates:
return
for row in duplicates:
model_id = row["model_id"]
conn.execute(
"""
DELETE FROM model_update_status
WHERE model_id = ?
AND rowid NOT IN (
SELECT rowid
FROM model_update_status
WHERE model_id = ?
ORDER BY
CASE WHEN last_checked_at IS NULL THEN 0 ELSE 1 END DESC,
last_checked_at DESC,
rowid DESC
LIMIT 1
)
""",
(model_id, model_id),
)
async def refresh_for_model_type(
self,
model_type: str,
scanner,
metadata_provider,
*,
force_refresh: bool = False,
target_model_ids: Optional[Sequence[int]] = None,
folder_path: Optional[str] = None,
) -> Dict[int, ModelUpdateRecord]:
"""Refresh update information for every model present in the cache."""
scanner.reset_cancellation()
normalized_targets = (
self._normalize_sequence(target_model_ids)
if target_model_ids is not None
else []
)
target_filter = normalized_targets or None
local_versions = await self._collect_local_versions(
scanner,
target_model_ids=target_filter,
folder_path=folder_path,
)
total_models = len(local_versions)
if total_models == 0:
if target_filter:
logger.info(
"No %s models matched requested ids %s while refreshing update metadata",
model_type,
target_filter,
)
else:
logger.info(
"No %s models found while refreshing update metadata", model_type
)
return {}
logger.info(
"Refreshing update metadata for %d %s models", total_models, model_type
)
# When filtering by folder, also collect the cross-folder version set
# so that versions already present in other folders are not reported
# as available updates. See issue #997.
all_local_versions: Optional[Dict[int, List[int]]] = None
if folder_path is not None:
all_local_versions = await self._collect_local_versions(
scanner,
target_model_ids=target_filter,
)
local_base_models = await self._collect_local_version_bases(
scanner,
target_model_ids=target_filter,
)
results: Dict[int, ModelUpdateRecord] = {}
prefetched: Dict[int, Mapping[Any, Any]] = {}
fetch_targets: List[int] = []
if metadata_provider and local_versions:
now = time.time()
async with self._lock:
for model_id in local_versions.keys():
existing = self._get_record(model_type, model_id)
if existing and existing.should_ignore_model and not force_refresh:
continue
if force_refresh or not existing or self._is_stale(existing, now):
fetch_targets.append(model_id)
if fetch_targets:
provider_name = (
metadata_provider.__class__.__name__
if metadata_provider is not None
else "unknown"
)
logger.info(
"Fetching remote metadata for %d %s models via bulk API using %s",
len(fetch_targets),
model_type,
provider_name,
)
try:
prefetched = await self._fetch_model_versions_bulk(
metadata_provider,
fetch_targets,
)
except NotImplementedError:
prefetched = {}
progress_interval = max(1, total_models // 10)
for index, (model_id, version_ids) in enumerate(
local_versions.items(), start=1
):
# Use cross-folder version IDs for is_in_library if available
all_vids: Sequence[int] = (
all_local_versions.get(model_id, [])
if all_local_versions is not None
else version_ids
)
record = await self._refresh_single_model(
model_type,
model_id,
version_ids,
metadata_provider,
force_refresh=force_refresh,
prefetched_response=prefetched.get(model_id),
all_local_version_ids=all_vids,
local_base_models=local_base_models,
)
if scanner.is_cancelled():
logger.info(f"{model_type.capitalize()} Update Service: Refresh cancelled by user")
return results
if record:
results[model_id] = record
if index % progress_interval == 0 or index == total_models:
logger.info(
"Refreshed update metadata for %d/%d %s models",
index,
total_models,
model_type,
)
logger.info(
"Completed update refresh for %d %s models; %d records stored",
total_models,
model_type,
len(results),
)
return results
async def refresh_single_model(
self,
model_type: str,
model_id: int,
scanner,
metadata_provider,
*,
force_refresh: bool = False,
) -> Optional[ModelUpdateRecord]:
"""Refresh update information for a specific model id."""
local_versions = await self._collect_local_versions(scanner)
version_ids = local_versions.get(model_id, [])
local_base_models = await self._collect_local_version_bases(scanner)
return await self._refresh_single_model(
model_type,
model_id,
version_ids,
metadata_provider,
force_refresh=force_refresh,
local_base_models=local_base_models,
)
async def update_in_library_versions(
self,
model_type: str,
model_id: int,
version_ids: Sequence[int],
*,
version_info: Optional[Mapping[str, Any]] = None,
) -> ModelUpdateRecord:
"""Persist a new set of in-library version identifiers."""
normalized_versions = self._normalize_sequence(version_ids)
async with self._lock:
existing = self._get_record(model_type, model_id)
record = self._merge_with_local_versions(
existing,
normalized_versions,
model_type=model_type,
model_id=model_id,
version_info=version_info,
)
self._upsert_record(record)
return record
async def set_should_ignore(
self, model_type: str, model_id: int, should_ignore: bool
) -> ModelUpdateRecord:
"""Toggle the ignore flag for a model."""
async with self._lock:
existing = self._get_record(model_type, model_id)
if existing:
record = ModelUpdateRecord(
model_type=existing.model_type,
model_id=existing.model_id,
versions=list(existing.versions),
last_checked_at=existing.last_checked_at,
should_ignore_model=should_ignore,
)
else:
record = ModelUpdateRecord(
model_type=model_type,
model_id=model_id,
versions=[],
last_checked_at=None,
should_ignore_model=should_ignore,
)
self._upsert_record(record)
return record
async def set_version_should_ignore(
self,
model_type: str,
model_id: int,
version_id: int,
should_ignore: bool,
) -> ModelUpdateRecord:
"""Toggle the ignore flag for an individual version."""
async with self._lock:
existing = self._get_record(model_type, model_id)
versions: List[ModelVersionRecord] = []
found = False
if existing:
for record_version in existing.versions:
if record_version.version_id == version_id:
versions.append(
replace(record_version, should_ignore=should_ignore)
)
found = True
else:
versions.append(record_version)
if not found:
versions.append(
ModelVersionRecord(
version_id=version_id,
name=None,
base_model=None,
released_at=None,
size_bytes=None,
preview_url=None,
is_in_library=False,
should_ignore=should_ignore,
sort_index=len(versions),
early_access_ends_at=None,
is_early_access=False,
)
)
record = ModelUpdateRecord(
model_type=existing.model_type if existing else model_type,
model_id=existing.model_id if existing else model_id,
versions=self._sorted_versions(versions),
last_checked_at=existing.last_checked_at if existing else None,
should_ignore_model=existing.should_ignore_model if existing else False,
)
self._upsert_record(record)
return record
async def get_record(self, model_type: str, model_id: int) -> Optional[ModelUpdateRecord]:
"""Return a cached record without triggering remote fetches."""
async with self._lock:
return self._get_record(model_type, model_id)
async def has_update(
self,
model_type: str,
model_id: int,
hide_early_access: bool = False,
hide_paid: bool = False,
) -> bool:
"""Determine if a model has updates pending."""
record = await self.get_record(model_type, model_id)
return (
record.has_update(
hide_early_access=hide_early_access, hide_paid=hide_paid
)
if record
else False
)
async def has_updates_bulk(
self,
model_type: str,
model_ids: Sequence[int],
hide_early_access: bool = False,
hide_paid: bool = False,
) -> Dict[int, bool]:
"""Return update availability for each model id in a single database pass."""
normalized_ids = self._normalize_sequence(model_ids)
if not normalized_ids:
return {}
async with self._lock:
records = self._get_records_bulk(model_type, normalized_ids)
return {
model_id: (
records[model_id].has_update(
hide_early_access=hide_early_access, hide_paid=hide_paid
)
if model_id in records
else False
)
for model_id in normalized_ids
}
async def get_records_bulk(
self,
model_type: str,
model_ids: Sequence[int],
) -> Dict[int, ModelUpdateRecord]:
"""Return cached update records for the requested models."""
normalized_ids = self._normalize_sequence(model_ids)
if not normalized_ids:
return {}
async with self._lock:
return self._get_records_bulk(model_type, normalized_ids)
async def get_price_alerts(
self, model_type: str, limit: int = 200
) -> List[Dict[str, Any]]:
"""Return the gated versions whose price currently crosses the threshold.
Reads the persisted ``price_alert_state`` rather than recomputing, so the
list matches the badges the user sees. Cheapest versions first, since that
is what a price watch is for.
"""
try:
normalized_limit = max(1, min(int(limit), 1000))
except (TypeError, ValueError):
normalized_limit = 200
with self._connect() as conn:
rows = conn.execute(
"""
SELECT v.model_id,
v.version_id,
v.name,
v.price_buzz,
v.list_price_buzz,
v.accepts_blue_buzz,
v.price_sale_ends_at,
v.price_checked_at,
v.is_in_library,
v.early_access_ends_at,
v.paid_access,
v.is_paid,
v.is_early_access,
s.model_type
FROM model_update_versions v
JOIN model_update_status s ON s.model_id = v.model_id
WHERE v.price_alert_state = 1
AND v.should_ignore = 0
AND s.should_ignore_model = 0
AND s.model_type = ?
ORDER BY v.price_buzz ASC, v.model_id ASC, v.version_id ASC
LIMIT ?
""",
(model_type, normalized_limit),
).fetchall()
return [
{
"modelId": int(row["model_id"]),
"modelType": row["model_type"],
"versionId": int(row["version_id"]),
"versionName": row["name"],
"priceBuzz": _normalize_int(row["price_buzz"]),
"listPriceBuzz": _normalize_int(row["list_price_buzz"]),
"acceptsBlueBuzz": bool(row["accepts_blue_buzz"]),
"priceSaleEndsAt": row["price_sale_ends_at"],
"priceCheckedAt": (
float(row["price_checked_at"])
if row["price_checked_at"] is not None
else None
),
"isInLibrary": bool(row["is_in_library"]),
"earlyAccessEndsAt": row["early_access_ends_at"],
"isPaid": bool(row["is_paid"]),
"isEarlyAccess": bool(row["is_early_access"]),
}
for row in rows
]
async def _refresh_single_model(
self,
model_type: str,
model_id: int,
local_versions: Sequence[int],
metadata_provider,
*,
force_refresh: bool = False,
prefetched_response: Optional[Mapping[str, Any]] = None,
all_local_version_ids: Optional[Sequence[int]] = None,
local_base_models: Optional[Mapping[int, str]] = None,
) -> Optional[ModelUpdateRecord]:
normalized_local = self._normalize_sequence(local_versions)
# When folder-filtering, this carries the cross-folder version set
# for is_in_library; otherwise it falls back to normalized_local.
normalized_all = (
self._normalize_sequence(all_local_version_ids)
if all_local_version_ids is not None
else normalized_local
)
now = time.time()
async with self._lock:
existing = self._get_record(model_type, model_id)
if existing and existing.should_ignore_model and not force_refresh:
record = self._merge_with_local_versions(
existing,
normalized_local,
all_local_version_ids=normalized_all,
)
self._upsert_record(record)
return record
should_fetch = force_refresh or not existing or self._is_stale(existing, now)
# release lock during network request
fetched_versions: List[ModelVersionRecord] | None = None
refresh_succeeded = False
fallback_attempted = False
fallback_error_message: Optional[str] = None
mark_model_as_ignored = False
response: Optional[Mapping[str, Any]] = None
if metadata_provider and should_fetch:
response = prefetched_response
if response is None:
fallback_attempted = True
try:
response = await metadata_provider.get_model_versions(model_id)
if response is not None:
await self._enrich_version_entries(
metadata_provider,
{model_id: response},
)
except RateLimitError:
raise
except ResourceNotFoundError as exc:
fallback_error_message = str(exc) or "resource not found"
mark_model_as_ignored = True
except Exception as exc: # pragma: no cover - defensive log
logger.warning(
"Failed to fetch versions for model %s (%s): %s",
model_id,
model_type,
exc,
)
fallback_error_message = str(exc)
if response is not None:
extracted = self._extract_versions(response)
if extracted is not None:
fetched_versions = extracted
refresh_succeeded = True
elif fallback_attempted and fallback_error_message is None:
fallback_error_message = "no versions returned"
elif fallback_attempted and fallback_error_message is None:
fallback_error_message = "no response"
# Optional price capture. Runs after the version list is known so only the
# models that actually carry a gate pay for a second (page) request, and
# stays outside the lock along with the other network I/O.
if (
refresh_succeeded
and isinstance(fetched_versions, list)
and fetched_versions
and metadata_provider is not None
and self._should_fetch_prices(fetched_versions, existing)
):
fetched_versions = await self._apply_model_prices(
metadata_provider, model_id, fetched_versions
)
if fallback_attempted:
if refresh_succeeded and isinstance(fetched_versions, list):
logger.info(
"Fetched metadata via single lookup for model %s (%s); received %d versions",
model_id,
model_type,
len(fetched_versions),
)
elif mark_model_as_ignored:
logger.info(
"Single lookup for model %s (%s) reported missing remote resource: %s",
model_id,
model_type,
fallback_error_message or "resource not found",
)
else:
logger.warning(
"Single lookup for model %s (%s) failed: %s",
model_id,
model_type,
fallback_error_message or "unknown error",
)
async with self._lock:
existing = self._get_record(model_type, model_id)
if existing and existing.should_ignore_model and not force_refresh:
record = self._merge_with_local_versions(
existing,
normalized_local,
all_local_version_ids=normalized_all,
)
self._upsert_record(record)
return record
if mark_model_as_ignored:
record = self._merge_with_local_versions(
existing,
normalized_local,
model_type=model_type,
model_id=model_id,
last_checked_at=now,
all_local_version_ids=normalized_all,
)
record = replace(record, should_ignore_model=True)
self._upsert_record(record)
logger.info(
"Marked model %s (%s) as ignored after remote resource was not found",
model_id,
model_type,
)
return record
if refresh_succeeded and isinstance(fetched_versions, list):
record = self._build_record_from_remote(
model_type,
model_id,
normalized_local,
fetched_versions,
existing,
now,
all_local_version_ids=normalized_all,
local_base_models=local_base_models,
)
else:
record = self._merge_with_local_versions(
existing,
normalized_local,
model_type=model_type,
model_id=model_id,
last_checked_at=existing.last_checked_at if existing else None,
all_local_version_ids=normalized_all,
)
self._upsert_record(record)
return record
async def _enrich_version_entries(
self,
metadata_provider,
responses_by_model_id: Dict[int, Mapping[Any, Any]],
) -> None:
"""Enrich version entries with ``usageControl`` via batch hash endpoint.
The model-level API does not include ``usageControl`` on version
entries. This method collects SHA256 hashes from every version's
primary model file, calls ``POST /api/v1/model-versions/by-hash``
(up to 100 hashes per request), and injects ``usageControl`` +
``earlyAccessEndsAt`` into each version entry dict in-place.
"""
if not metadata_provider or not responses_by_model_id:
return
hashes_by_version: Dict[int, str] = {}
for response in responses_by_model_id.values():
hashes_by_version.update(
self._collect_hashes_from_response(response)
)
if not hashes_by_version:
return
version_ids_by_hash: Dict[str, List[int]] = {}
for version_id, sha256 in hashes_by_version.items():
version_ids_by_hash.setdefault(sha256, []).append(version_id)
all_hashes = list(version_ids_by_hash.keys())
BATCH_SIZE = 100
enrichment: Dict[int, Dict[str, Any]] = {}
try:
for start in range(0, len(all_hashes), BATCH_SIZE):
batch = all_hashes[start : start + BATCH_SIZE]
try:
enriched = await metadata_provider.get_model_versions_by_hashes(
batch
)
except NotImplementedError:
return
except RateLimitError:
raise
except Exception:
continue
if not enriched:
continue
for entry in enriched:
if not isinstance(entry, dict):
continue
version_id = entry.get("id")
if version_id is None:
continue
enrichment[version_id] = {
"usageControl": _normalize_string(
entry.get("usageControl")
),
"earlyAccessEndsAt": _normalize_string(
entry.get("earlyAccessEndsAt")
),
"paidAccess": entry.get("paidAccess"),
}
except RateLimitError:
raise
if not enrichment:
return
for response in responses_by_model_id.values():
versions = response.get("modelVersions")
if not isinstance(versions, list):
continue
for version in versions:
if not isinstance(version, dict):
continue
version_id = version.get("id")
if version_id not in enrichment:
continue
extra = enrichment[version_id]
if extra.get("usageControl") and not version.get("usageControl"):
version["usageControl"] = extra["usageControl"]
if extra.get("earlyAccessEndsAt") and not version.get(
"earlyAccessEndsAt"
):
version["earlyAccessEndsAt"] = extra["earlyAccessEndsAt"]
# Only backfill when the model-level response carries no *active*
# paidAccess signal: a present-but-empty DTO (e.g.
# {"permanent": false, "endsAt": null}) would otherwise block
# the authoritative by-hash data.
extra_paid = ModelUpdateService._normalize_paid_access(
extra.get("paidAccess")
)
if extra_paid and not ModelUpdateService._normalize_paid_access(
version.get("paidAccess")
):
version["paidAccess"] = extra["paidAccess"]
@staticmethod
def _collect_hashes_from_response(response: Mapping[str, Any]) -> Dict[int, str]:
"""Extract ``{version_id: sha256}`` from a model-level API response.
Returns an empty dict if the response structure is unexpected.
"""
result: Dict[int, str] = {}
versions = response.get("modelVersions")
if not isinstance(versions, list):
return result
for entry in versions:
if not isinstance(entry, dict):
continue
version_id = _normalize_int(entry.get("id"))
if version_id is None:
continue
sha256 = ModelUpdateService._extract_sha256_from_version_entry(entry)
if sha256:
result[version_id] = sha256
return result
@staticmethod
def _extract_sha256_from_version_entry(entry: Mapping[str, Any]) -> Optional[str]:
"""Return the SHA256 hash from the primary model file of a version entry."""
files = entry.get("files")
if not isinstance(files, list):
return None
for file_info in files:
if not isinstance(file_info, dict):
continue
if file_info.get("type") != "Model":
continue
primary = file_info.get("primary")
if primary is not True and str(primary).strip().lower() != "true":
continue
hashes = file_info.get("hashes")
if isinstance(hashes, dict):
sha256 = hashes.get("SHA256")
if sha256:
return sha256
return None
async def _fetch_model_versions_bulk(
self,
metadata_provider,
model_ids: Sequence[int],
) -> Dict[int, Mapping[Any, Any]]:
"""Fetch model metadata in batches of up to 100 ids."""
BATCH_SIZE = 100
normalized = self._normalize_sequence(model_ids)
provider = metadata_provider
if not normalized or provider is None:
return {}
aggregated: Dict[int, Mapping[Any, Any]] = {}
total_ids = len(normalized)
total_batches = (total_ids + BATCH_SIZE - 1) // BATCH_SIZE
provider_name = provider.__class__.__name__
for batch_index, start in enumerate(range(0, total_ids, BATCH_SIZE), start=1):
chunk = normalized[start : start + BATCH_SIZE]
logger.info(
"Requesting bulk metadata for %d models (batch %d/%d) from %s",
len(chunk),
batch_index,
total_batches,
provider_name,
)
try:
response = await provider.get_model_versions_bulk(chunk)
except RateLimitError:
raise
if response is None:
continue
if not isinstance(response, Mapping):
logger.debug(
"Unexpected bulk response type %s from provider %s", type(response), metadata_provider
)
continue
for key, value in response.items():
normalized_key = _normalize_int(key)
if normalized_key is None:
continue
if isinstance(value, Mapping):
aggregated[normalized_key] = value
logger.info(
"Completed bulk metadata fetch for %d models using %s",
len(aggregated),
provider_name,
)
await self._enrich_version_entries(metadata_provider, aggregated)
return aggregated
@staticmethod
def _iter_local_civitai_items(
cache,
*,
target_set: Optional[set[int]] = None,
normalized_folder: Optional[str] = None,
) -> Iterator[tuple[int, int, Any]]:
"""Yield ``(modelId, versionId, base_model)`` for each scannable item."""
if not cache or not getattr(cache, "raw_data", None):
return
for item in cache.raw_data:
# Apply folder filter first (cheapest check)
if normalized_folder is not None:
if not isinstance(item, dict):
continue
item_folder = (item.get("folder") or "").replace("\\", "/").strip("/")
if item_folder != normalized_folder and not item_folder.startswith(normalized_folder + "/"):
continue
civitai = item.get("civitai") if isinstance(item, dict) else None
if not isinstance(civitai, dict):
continue
model_id = _normalize_int(civitai.get("modelId"))
version_id = _normalize_int(civitai.get("id"))
if model_id is None or version_id is None:
continue
if target_set is not None and model_id not in target_set:
continue
yield model_id, version_id, item.get("base_model")
def _prepare_collection_filters(
self,
target_model_ids: Optional[Sequence[int]],
folder_path: Optional[str],
) -> tuple[Optional[set[int]], Optional[str]]:
target_set: Optional[set[int]] = None
if target_model_ids:
target_set = set(target_model_ids)
normalized_folder = None
if folder_path is not None:
normalized_folder = folder_path.replace("\\", "/").strip("/")
return target_set, normalized_folder
async def _collect_local_versions(
self,
scanner,
*,
target_model_ids: Optional[Sequence[int]] = None,
folder_path: Optional[str] = None,
) -> Dict[int, List[int]]:
cache = await scanner.get_cached_data()
mapping: Dict[int, set[int]] = {}
target_set, normalized_folder = self._prepare_collection_filters(
target_model_ids, folder_path
)
if target_model_ids and not target_set:
return {}
for model_id, version_id, _base_model in self._iter_local_civitai_items(
cache, target_set=target_set, normalized_folder=normalized_folder
):
mapping.setdefault(model_id, set()).add(version_id)
return {model_id: sorted(ids) for model_id, ids in mapping.items()}
async def _collect_local_version_bases(
self,
scanner,
*,
target_model_ids: Optional[Sequence[int]] = None,
) -> Dict[int, str]:
"""Map version id -> base model from cache items.
Deliberately unfiltered by folder: synthesized in-library entries must
carry a base regardless of which folder triggered the refresh.
"""
cache = await scanner.get_cached_data()
bases: Dict[int, str] = {}
target_set, _normalized_folder = self._prepare_collection_filters(
target_model_ids, None
)
if target_model_ids and not target_set:
return {}
for _model_id, version_id, base_model in self._iter_local_civitai_items(
cache, target_set=target_set
):
normalized_base = _normalize_string(base_model)
if normalized_base:
bases[version_id] = normalized_base
return bases
def _merge_with_local_versions(
self,
existing: Optional[ModelUpdateRecord],
normalized_local: Sequence[int],
*,
all_local_version_ids: Optional[Sequence[int]] = None,
model_type: Optional[str] = None,
model_id: Optional[int] = None,
last_checked_at: Optional[float] = None,
version_info: Optional[Mapping[str, Any]] = None,
) -> ModelUpdateRecord:
local_set = set(normalized_local)
# When folder-filtering, also consider versions in other folders
# as in-library so they are not reported as available updates.
effective_local_set: set[int] = (
local_set | set(all_local_version_ids)
if all_local_version_ids is not None
else local_set
)
versions: List[ModelVersionRecord] = []
ignore_map: Dict[int, bool] = {}
if existing:
model_type = existing.model_type
model_id = existing.model_id
last_checked_at = existing.last_checked_at if last_checked_at is None else last_checked_at
ignore_map = {version.version_id: version.should_ignore for version in existing.versions}
for version in existing.versions:
versions.append(
replace(
version,
is_in_library=version.version_id in effective_local_set,
)
)
elif model_type is None or model_id is None:
raise ValueError("model_type and model_id are required when creating a new record")
seen_ids = {version.version_id for version in versions}
for missing_id in sorted(local_set - seen_ids):
new_version: Optional[ModelVersionRecord] = None
if version_info and _normalize_int(version_info.get("id")) == missing_id:
new_version = self._extract_single_version(version_info, index=len(versions))
if new_version:
versions.append(replace(new_version, is_in_library=True))
else:
versions.append(
ModelVersionRecord(
version_id=missing_id,
name=None,
base_model=None,
released_at=None,
size_bytes=None,
preview_url=None,
is_in_library=True,
should_ignore=ignore_map.get(missing_id, False),
sort_index=len(versions),
early_access_ends_at=None,
is_early_access=False,
)
)
return ModelUpdateRecord(
model_type=model_type,
model_id=model_id,
versions=self._sorted_versions(versions),
last_checked_at=last_checked_at,
should_ignore_model=existing.should_ignore_model if existing else False,
)
def _build_record_from_remote(
self,
model_type: str,
model_id: int,
local_versions: Sequence[int],
remote_versions: Sequence[ModelVersionRecord],
existing: Optional[ModelUpdateRecord],
timestamp: float,
*,
all_local_version_ids: Optional[Sequence[int]] = None,
local_base_models: Optional[Mapping[int, str]] = None,
) -> ModelUpdateRecord:
local_set = set(local_versions)
# When folder-filtering, also consider versions in other folders
# as in-library so they are not reported as available updates.
effective_local_set: set[int] = (
local_set | set(all_local_version_ids)
if all_local_version_ids is not None
else local_set
)
ignore_map = {version.version_id: version.should_ignore for version in existing.versions} if existing else {}
preview_map = {version.version_id: version.preview_url for version in existing.versions} if existing else {}
file_count_map = {version.version_id: version.file_count for version in existing.versions} if existing else {}
sort_map = {version.version_id: version.sort_index for version in existing.versions} if existing else {}
existing_map = {version.version_id: version for version in existing.versions} if existing else {}
versions: List[ModelVersionRecord] = []
events: List[Dict[str, Any]] = []
seen_ids: set[int] = set()
lapse_timestamp = _format_utc_timestamp(timestamp)
for index, remote_version in enumerate(remote_versions):
version_id = remote_version.version_id
seen_ids.add(version_id)
existing_version = existing_map.get(version_id)
is_gated = self._has_structural_gate(remote_version)
was_gated = (
self._has_structural_gate(existing_version)
if existing_version is not None
else None
)
should_ignore = ignore_map.get(version_id, remote_version.should_ignore)
# `gate_lapsed_at` survives refreshes so a version that became free keeps
# its "free since" marker, and is cleared if the gate comes back.
if is_gated:
gate_lapsed_at = None
elif was_gated:
gate_lapsed_at = lapse_timestamp
elif existing_version is not None:
gate_lapsed_at = existing_version.gate_lapsed_at
else:
gate_lapsed_at = None
price_fields = self._price_fields_for(
remote_version, existing_version, is_gated=is_gated
)
price_fields["price_alert_state"] = self._evaluate_price_alert(
price_fields.get("price_buzz")
)
if (
existing_version is not None
and not should_ignore
and was_gated != is_gated
):
events.append(
{
"versionId": version_id,
"kind": "new_gate" if is_gated else "became_free",
"versionName": remote_version.name,
"isInLibrary": version_id in effective_local_set,
}
)
# Only the edge fires: a version that keeps sitting under the threshold
# must not re-announce itself on every refresh.
if (
existing_version is not None
and not should_ignore
and price_fields["price_alert_state"]
and not existing_version.price_alert_state
):
events.append(
{
"versionId": version_id,
"kind": "price_drop",
"versionName": remote_version.name,
"isInLibrary": version_id in effective_local_set,
"priceBuzz": price_fields.get("price_buzz"),
}
)
versions.append(
ModelVersionRecord(
version_id=version_id,
name=remote_version.name,
base_model=remote_version.base_model,
released_at=remote_version.released_at,
size_bytes=remote_version.size_bytes,
preview_url=remote_version.preview_url or preview_map.get(version_id),
is_in_library=version_id in effective_local_set,
should_ignore=should_ignore,
sort_index=sort_map.get(version_id, index),
early_access_ends_at=remote_version.early_access_ends_at,
is_early_access=remote_version.is_early_access,
usage_control=remote_version.usage_control,
paid_access=remote_version.paid_access,
is_paid=remote_version.is_paid,
file_count=(
remote_version.file_count
if remote_version.file_count is not None
else file_count_map.get(version_id)
),
gate_lapsed_at=gate_lapsed_at,
**price_fields,
)
)
missing_local = local_set - seen_ids
if missing_local:
item_base_models = local_base_models or {}
for version_id in sorted(missing_local):
existing_version = existing_map.get(version_id)
if existing_version:
versions.append(
replace(
existing_version,
is_in_library=True,
)
)
else:
versions.append(
ModelVersionRecord(
version_id=version_id,
name=None,
base_model=item_base_models.get(version_id),
released_at=None,
size_bytes=None,
preview_url=None,
is_in_library=True,
should_ignore=ignore_map.get(version_id, False),
sort_index=len(versions),
early_access_ends_at=None,
is_early_access=False,
)
)
return ModelUpdateRecord(
model_type=model_type,
model_id=model_id,
versions=self._sorted_versions(versions),
last_checked_at=timestamp,
should_ignore_model=existing.should_ignore_model if existing else False,
events=events,
)
@staticmethod
def _has_structural_gate(version: Optional[ModelVersionRecord]) -> bool:
"""True when a version carries a paid gate, ignoring the clock.
Transition detection uses this rather than a time-based check so a gate
that lapsed since the previous refresh still reads as "was gated" — the
public API reports no ``paidAccess`` at all once a gate expires, so the
time-based helper alone would silently see "free before, free now".
"""
if version is None:
return False
return bool(version.paid_access) or version.is_paid or version.is_early_access
@staticmethod
def _price_fields_for(
remote_version: ModelVersionRecord,
existing_version: Optional[ModelVersionRecord],
*,
is_gated: bool,
) -> Dict[str, Any]:
"""Resolve the price columns for a refreshed version.
A refresh only carries prices when a price fetch actually ran (recognizable
by ``price_checked_at``), so otherwise the previously stored price is kept —
but only while the version is still gated: once it is free the stored
numbers are stale and must not keep firing alerts.
"""
if not is_gated:
return {
"price_buzz": None,
"list_price_buzz": None,
"generation_price_buzz": None,
"accepts_blue_buzz": False,
"price_sale_ends_at": None,
"price_checked_at": None,
"price_alert_state": False,
}
if remote_version.price_checked_at is not None:
source = remote_version
else:
source = existing_version or remote_version
return {
"price_buzz": source.price_buzz,
"list_price_buzz": source.list_price_buzz,
"generation_price_buzz": source.generation_price_buzz,
"accepts_blue_buzz": source.accepts_blue_buzz,
"price_sale_ends_at": source.price_sale_ends_at,
"price_checked_at": source.price_checked_at,
"price_alert_state": source.price_alert_state,
}
def _price_tracking_enabled(self) -> bool:
"""Whether the optional (page-fetching) price capture is switched on."""
try:
return bool(self._settings.get("price_tracking_enabled", False))
except Exception:
return False
def _price_alert_threshold_buzz(self) -> int:
"""Buzz threshold for a price alert; 0 means "free only"."""
try:
value = int(self._settings.get("price_alert_threshold_buzz", 0) or 0)
except (TypeError, ValueError):
return 0
return max(0, value)
def _evaluate_price_alert(self, price_buzz: Optional[int]) -> bool:
"""Whether a version's price crosses the user's alert threshold.
Free versions never carry a price, so a threshold of 0 alerts only via the
``became_free`` event — which is exactly the "or free" half of the request.
"""
if not self._price_tracking_enabled():
return False
if price_buzz is None:
return False
return price_buzz <= self._price_alert_threshold_buzz()
def _price_check_ttl_seconds(self) -> float:
try:
hours = float(self._settings.get("price_check_ttl_hours", 24))
except (TypeError, ValueError):
hours = 24.0
if hours <= 0:
hours = 24.0
return hours * 3600.0
def _should_fetch_prices(
self,
remote_versions: Sequence[ModelVersionRecord],
existing: Optional[ModelUpdateRecord],
) -> bool:
"""Whether this model needs a price fetch this round.
Only gated versions are considered, so an unmodified free library pays
nothing. A stored price is refreshed once its TTL lapses, and immediately
when the gate itself changed (a new end date or sale window is a reason to
believe the price moved).
"""
if not self._price_tracking_enabled():
return False
existing_map = (
{version.version_id: version for version in existing.versions}
if existing
else {}
)
now = time.time()
ttl = self._price_check_ttl_seconds()
for remote_version in remote_versions:
if not self._has_structural_gate(remote_version):
continue
stored = existing_map.get(remote_version.version_id)
if stored is None:
return True
if stored.price_checked_at is None:
return True
if (now - stored.price_checked_at) >= ttl:
return True
if (stored.paid_access or None) != (remote_version.paid_access or None):
return True
return False
async def _apply_model_prices(
self,
metadata_provider,
model_id: int,
versions: Sequence[ModelVersionRecord],
) -> List[ModelVersionRecord]:
"""Attach freshly fetched prices to the versions that have them.
Never raises for a provider problem: price tracking is a convenience, and
an unreadable page (or a provider that has no prices at all) must leave
the update check exactly as it was.
"""
getter = getattr(metadata_provider, "get_model_prices", None)
if not callable(getter):
return list(versions)
try:
prices = await getter(model_id)
except RateLimitError:
raise
except Exception as exc: # pragma: no cover - defensive
logger.debug("Price fetch failed for model %s: %s", model_id, exc)
return list(versions)
if not isinstance(prices, Mapping) or not prices:
return list(versions)
checked_at = time.time()
enriched: List[ModelVersionRecord] = []
for version in versions:
fields = prices.get(version.version_id)
if not isinstance(fields, Mapping):
# Not priced (or not understood): leave the stored values alone so
# the next refresh retries instead of recording a blank price.
enriched.append(version)
continue
recognized = {
key: value
for key, value in fields.items()
if key in _PRICE_FIELD_NAMES
}
if not recognized:
enriched.append(version)
continue
enriched.append(
replace(version, price_checked_at=checked_at, **recognized)
)
return enriched
def _sorted_versions(self, versions: Sequence[ModelVersionRecord]) -> List[ModelVersionRecord]:
ordered = sorted(versions, key=lambda version: (version.sort_index, version.version_id))
normalized: List[ModelVersionRecord] = []
for index, version in enumerate(ordered):
normalized.append(replace(version, sort_index=index))
return normalized
def _is_stale(self, record: ModelUpdateRecord, now: float) -> bool:
if record.last_checked_at is None:
return True
return (now - record.last_checked_at) >= self._ttl_seconds
def _normalize_sequence(self, values: Sequence[int]) -> List[int]:
normalized = [
item
for item in (_normalize_int(value) for value in values)
if item is not None
]
return sorted(dict.fromkeys(normalized))
def _extract_versions(self, response) -> Optional[List[ModelVersionRecord]]:
if not isinstance(response, Mapping):
return None
versions = response.get("modelVersions")
if versions is None:
return []
if not isinstance(versions, Iterable):
return None
extracted: List[ModelVersionRecord] = []
for index, entry in enumerate(versions):
version_record = self._extract_single_version(entry, index)
if version_record:
extracted.append(version_record)
return extracted
def _extract_single_version(
self, entry: Any, index: int = 0
) -> Optional[ModelVersionRecord]:
"""Convert a raw metadata entry into a structured record."""
if not isinstance(entry, Mapping):
return None
version_id = _normalize_int(entry.get("id"))
if version_id is None:
return None
name = _normalize_string(entry.get("name"))
base_model = _normalize_string(entry.get("baseModel"))
released_at = _normalize_string(entry.get("publishedAt") or entry.get("createdAt"))
size_bytes = self._extract_size_bytes(entry.get("files"))
file_count = self._extract_file_count(entry.get("files"))
preview_url = self._extract_preview_url(entry.get("images"))
early_access_ends_at = _normalize_string(entry.get("earlyAccessEndsAt"))
# Check availability field from bulk API for basic EA detection
availability = _normalize_string(entry.get("availability"))
is_early_access = availability == "EarlyAccess"
usage_control = _normalize_string(entry.get("usageControl"))
# CivitAI's paidAccess DTO ({"permanent": bool, "endsAt": ISO|null})
# gates versions behind a paid tier while availability stays "Public".
# A non-null DTO from the public API is always an ACTIVE gate: lapsed
# (tombstone) gates come back as null. That includes the timed gate whose
# end is not recorded yet, {"permanent": false, "endsAt": null}.
paid_access = self._normalize_paid_access(entry.get("paidAccess"))
paid_access_json = json.dumps(paid_access) if paid_access else None
is_paid = bool(paid_access.get("permanent")) if paid_access else False
if early_access_ends_at is None and paid_access and paid_access.get("endsAt"):
early_access_ends_at = _normalize_string(paid_access.get("endsAt"))
# Only timed gates are early access; permanent paid versions are not
# (consumers filter them via is_paid), so the stored flag stays accurate.
if not is_early_access and paid_access and not paid_access.get("permanent"):
is_early_access = True
return ModelVersionRecord(
version_id=version_id,
name=name,
base_model=base_model,
released_at=released_at,
size_bytes=size_bytes,
preview_url=preview_url,
is_in_library=False,
should_ignore=False,
early_access_ends_at=early_access_ends_at,
sort_index=index,
is_early_access=is_early_access,
usage_control=usage_control,
paid_access=paid_access_json,
is_paid=is_paid,
file_count=file_count,
)
@staticmethod
def _normalize_paid_access(value) -> Optional[Dict[str, Any]]:
"""Normalize a CivitAI ``paidAccess`` DTO into a mapping.
Accepts a dict, None, or a JSON string (as carried by the by-hash
enrichment path) and returns ``{"permanent": bool, "endsAt": str|None}``
or None when the input carries no paid-access signal.
Delegates to :mod:`py.utils.paid_access` so the update service and the
download gate cannot disagree about what counts as a gate.
"""
return _normalize_paid_access_payload(value)
@staticmethod
def _extract_file_count(files) -> Optional[int]:
"""Count downloadable weight files in a version entry's ``files`` list.
Returns None when the payload carries no files array (unknown), so
callers can distinguish "no weight files" from "no data".
"""
if not isinstance(files, list):
return None
count = 0
for entry in files:
if not isinstance(entry, Mapping):
continue
entry_type = entry.get("type")
if isinstance(entry_type, str) and entry_type in MODEL_WEIGHT_FILE_TYPES:
count += 1
return count
def _extract_size_bytes(self, files) -> Optional[int]:
if not isinstance(files, Iterable):
return None
def parse_size(entry: Mapping[str, Any]) -> Optional[int]:
size_kb = entry.get("sizeKB")
if size_kb is None:
return None
try:
return int(float(size_kb) * 1024)
except (TypeError, ValueError):
return None
preferred_size: Optional[int] = None
fallback_size: Optional[int] = None
for entry in files:
if not isinstance(entry, Mapping):
continue
size_bytes = parse_size(entry)
if size_bytes is None:
continue
entry_type = entry.get("type")
is_model_type = isinstance(entry_type, str) and entry_type.lower() == "model"
primary_flag = entry.get("primary")
is_primary = primary_flag is True or (
isinstance(primary_flag, str) and primary_flag.strip().lower() == "true"
)
if is_model_type and is_primary:
preferred_size = size_bytes
break
if fallback_size is None:
fallback_size = size_bytes
return preferred_size if preferred_size is not None else fallback_size
def _extract_preview_url(self, images) -> Optional[str]:
if not isinstance(images, Iterable):
return None
candidates = [entry for entry in images if isinstance(entry, Mapping)]
if not candidates:
return None
blur_mature_content = True
mature_threshold = resolve_mature_threshold({"mature_blur_level": "R"})
settings = getattr(self, "_settings", None)
if settings is not None and hasattr(settings, "get"):
try:
blur_mature_content = bool(settings.get("blur_mature_content", True))
mature_threshold = resolve_mature_threshold(
{"mature_blur_level": settings.get("mature_blur_level", "R")}
)
except Exception: # pragma: no cover - defensive guard
blur_mature_content = True
mature_threshold = resolve_mature_threshold({"mature_blur_level": "R"})
selected, _ = select_preview_media(
candidates,
blur_mature_content=blur_mature_content,
mature_threshold=mature_threshold,
)
if not selected:
return None
url = selected.get("url")
if not isinstance(url, str) or not url:
return None
media_type = selected.get("type")
if not isinstance(media_type, str):
media_type = None
rewritten, _ = rewrite_preview_url(url, media_type)
return rewritten or url
def _get_record(self, model_type: str, model_id: int) -> Optional[ModelUpdateRecord]:
records = self._get_records_bulk(model_type, [model_id])
return records.get(model_id)
def _get_records_bulk(
self,
model_type: str,
model_ids: Sequence[int],
) -> Dict[int, ModelUpdateRecord]:
if not model_ids:
return {}
ids = list(model_ids)
status_rows: list[sqlite3.Row] = []
version_rows: list[sqlite3.Row] = []
with self._connect() as conn:
for start in range(0, len(ids), self._SQLITE_MAX_VARIABLES):
chunk = tuple(ids[start : start + self._SQLITE_MAX_VARIABLES])
placeholders = ",".join("?" for _ in chunk)
chunk_status = conn.execute(
f"""
SELECT model_id, model_type, last_checked_at, should_ignore_model
FROM model_update_status
WHERE model_id IN ({placeholders})
""",
chunk,
).fetchall()
status_rows.extend(chunk_status)
chunk_versions = conn.execute(
f"""
SELECT model_id, version_id, sort_index, name, base_model, released_at,
size_bytes, preview_url, is_in_library, should_ignore, early_access_ends_at,
is_early_access, usage_control, paid_access, is_paid, file_count,
gate_lapsed_at, price_buzz, list_price_buzz, generation_price_buzz,
accepts_blue_buzz, price_sale_ends_at, price_checked_at, price_alert_state
FROM model_update_versions
WHERE model_id IN ({placeholders})
ORDER BY model_id ASC, sort_index ASC, version_id ASC
""",
chunk,
).fetchall()
version_rows.extend(chunk_versions)
if not status_rows:
return {}
versions_by_model: Dict[int, List[ModelVersionRecord]] = {}
for row in version_rows:
model_id = int(row["model_id"])
versions_by_model.setdefault(model_id, []).append(
ModelVersionRecord(
version_id=int(row["version_id"]),
name=row["name"],
base_model=row["base_model"],
released_at=row["released_at"],
size_bytes=_normalize_int(row["size_bytes"]),
preview_url=row["preview_url"],
is_in_library=bool(row["is_in_library"]),
should_ignore=bool(row["should_ignore"]),
early_access_ends_at=row["early_access_ends_at"],
sort_index=_normalize_int(row["sort_index"]) or 0,
is_early_access=bool(row["is_early_access"]),
usage_control=row["usage_control"],
paid_access=row["paid_access"],
is_paid=bool(row["is_paid"]),
file_count=_normalize_int(row["file_count"]),
gate_lapsed_at=row["gate_lapsed_at"],
price_buzz=_normalize_int(row["price_buzz"]),
list_price_buzz=_normalize_int(row["list_price_buzz"]),
generation_price_buzz=_normalize_int(row["generation_price_buzz"]),
accepts_blue_buzz=bool(row["accepts_blue_buzz"]),
price_sale_ends_at=row["price_sale_ends_at"],
price_checked_at=(
float(row["price_checked_at"])
if row["price_checked_at"] is not None
else None
),
price_alert_state=bool(row["price_alert_state"]),
)
)
records: Dict[int, ModelUpdateRecord] = {}
for status in status_rows:
model_id = int(status["model_id"])
stored_type = status["model_type"]
if stored_type and stored_type != model_type:
logger.debug(
"Model id %s requested as %s but stored as %s",
model_id,
model_type,
stored_type,
)
record = ModelUpdateRecord(
model_type=stored_type or model_type,
model_id=model_id,
versions=self._sorted_versions(versions_by_model.get(model_id, [])),
last_checked_at=status["last_checked_at"],
should_ignore_model=bool(status["should_ignore_model"]),
)
records[model_id] = record
return records
def _upsert_record(self, record: ModelUpdateRecord) -> None:
payload = (
record.model_id,
record.model_type,
record.last_checked_at,
1 if record.should_ignore_model else 0,
)
with self._connect() as conn:
conn.execute(
"""
INSERT INTO model_update_status (
model_id, model_type, last_checked_at, should_ignore_model
) VALUES (?, ?, ?, ?)
ON CONFLICT(model_id) DO UPDATE SET
model_type = excluded.model_type,
last_checked_at = excluded.last_checked_at,
should_ignore_model = excluded.should_ignore_model
""",
payload,
)
conn.execute(
"DELETE FROM model_update_versions WHERE model_id = ?",
(record.model_id,),
)
for version in record.versions:
paid_access_value = (
version.paid_access
if version.paid_access is None
or isinstance(version.paid_access, str)
else json.dumps(version.paid_access)
)
conn.execute(
"""
INSERT INTO model_update_versions (
version_id, model_id, sort_index, name, base_model, released_at,
size_bytes, preview_url, is_in_library, should_ignore, early_access_ends_at,
is_early_access, usage_control, paid_access, is_paid, file_count,
gate_lapsed_at, price_buzz, list_price_buzz, generation_price_buzz,
accepts_blue_buzz, price_sale_ends_at, price_checked_at, price_alert_state
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
""",
(
version.version_id,
record.model_id,
version.sort_index,
version.name,
version.base_model,
version.released_at,
version.size_bytes,
version.preview_url,
1 if version.is_in_library else 0,
1 if version.should_ignore else 0,
version.early_access_ends_at,
1 if version.is_early_access else 0,
version.usage_control,
paid_access_value,
1 if version.is_paid else 0,
version.file_count,
version.gate_lapsed_at,
version.price_buzz,
version.list_price_buzz,
version.generation_price_buzz,
1 if version.accepts_blue_buzz else 0,
version.price_sale_ends_at,
version.price_checked_at,
1 if version.price_alert_state else 0,
),
)
conn.commit()
Reference in New Issue View Git Blame Copy Permalink