661 lines
28 KiB
Python
661 lines
28 KiB
Python
"""
|
|
Cache utilities for SageUtils.
|
|
Handles persistent storage and retrieval of model metadata, hashes, and info.
|
|
"""
|
|
|
|
import json
|
|
import pathlib
|
|
import hashlib
|
|
import datetime
|
|
import tempfile
|
|
import copy
|
|
import logging
|
|
import os
|
|
from typing import Any, Dict, Optional, List
|
|
|
|
from .path_manager import path_manager, file_manager
|
|
|
|
def str_to_bool(value: Any) -> bool:
|
|
if isinstance(value, bool):
|
|
return value
|
|
if isinstance(value, str):
|
|
value = value.lower()
|
|
if value in {'true', '1', 'yes'}:
|
|
return True
|
|
elif value in {'false', '0', 'no'}:
|
|
return False
|
|
raise ValueError(f"Cannot convert {value} to boolean.")
|
|
|
|
def bool_to_str(value: Any) -> str:
|
|
if isinstance(value, bool):
|
|
return "true" if value else "false"
|
|
elif isinstance(value, str):
|
|
value = value.lower()
|
|
if value in {'true', '1', 'yes'}:
|
|
return "true"
|
|
elif value in {'false', '0', 'no'}:
|
|
return "false"
|
|
raise ValueError(f"Cannot convert {value} to string representation of boolean.")
|
|
|
|
|
|
class SageCache:
|
|
"""
|
|
Persistent cache for model metadata, hashes, and info.
|
|
"""
|
|
def __init__(self):
|
|
# Use centralized path management
|
|
self.main_path = path_manager.get_user_file_path("sage_cache.json")
|
|
self.info_path = path_manager.get_user_file_path("sage_cache_info.json")
|
|
self.hash_path = path_manager.get_user_file_path("sage_cache_hash.json")
|
|
self.ollama_models_path = path_manager.get_user_file_path("sage_cache_ollama.json")
|
|
|
|
self.data: Dict[str, Any] = {}
|
|
self.hash: Dict[str, str] = {}
|
|
self.info: Dict[str, Any] = {}
|
|
self.ollama_models: Dict[str, Any] = {}
|
|
self.last_hash: Dict[str, str] = {}
|
|
self.last_info: Dict[str, Any] = {}
|
|
self.last_ollama_models: Dict[str, Any] = {}
|
|
self.num_of_backups_to_keep = 7
|
|
self.backup_counter = 0
|
|
|
|
# Batch mode attributes for deferred saves/backups
|
|
self.batch_mode = False
|
|
self.pending_changes = 0
|
|
self.batch_start_changes = 0
|
|
self.save_count_since_backup = 0
|
|
self.last_backup_time = datetime.datetime.now()
|
|
|
|
# Configuration thresholds
|
|
self.save_threshold = 10 # Save after N changes in batch mode (safety)
|
|
self.backup_threshold = 50 # Backup after N saves
|
|
self.backup_interval_seconds = 300 # 5 minutes between backups
|
|
|
|
# Backup manifest for fast comparison
|
|
self.backup_manifest_path = path_manager.get_backup_file_path("backup_manifest.json")
|
|
self.backup_manifest: Dict[str, Dict[str, Any]] = {}
|
|
self._load_backup_manifest()
|
|
|
|
self.prune_all_backups()
|
|
|
|
def _load_backup_manifest(self) -> None:
|
|
"""Load the backup manifest from disk."""
|
|
if self.backup_manifest_path.exists():
|
|
try:
|
|
with self.backup_manifest_path.open('r') as f:
|
|
self.backup_manifest = json.load(f)
|
|
logging.debug(f"Loaded backup manifest with {len(self.backup_manifest)} entries")
|
|
except Exception as e:
|
|
logging.warning(f"Failed to load backup manifest: {e}")
|
|
self.backup_manifest = {}
|
|
else:
|
|
self.backup_manifest = {}
|
|
|
|
def _save_backup_manifest(self) -> None:
|
|
"""Save the backup manifest to disk."""
|
|
try:
|
|
with self.backup_manifest_path.open('w') as f:
|
|
json.dump(self.backup_manifest, f, indent=2)
|
|
except Exception as e:
|
|
logging.warning(f"Failed to save backup manifest: {e}")
|
|
|
|
def _update_manifest_for_backup(self, backup_path: pathlib.Path, data: Any) -> None:
|
|
"""Update manifest entry for a backup file."""
|
|
try:
|
|
data_json = json.dumps(data, separators=(",", ":"), sort_keys=True)
|
|
entry_count = len(data) if isinstance(data, dict) else 0
|
|
file_size = len(data_json)
|
|
content_hash = hashlib.sha256(data_json.encode("utf-8")).hexdigest()
|
|
|
|
self.backup_manifest[backup_path.name] = {
|
|
"timestamp": datetime.datetime.now().isoformat(),
|
|
"entry_count": entry_count,
|
|
"file_size": file_size,
|
|
"content_hash": content_hash
|
|
}
|
|
self._save_backup_manifest()
|
|
except Exception as e:
|
|
logging.warning(f"Failed to update manifest for {backup_path.name}: {e}")
|
|
|
|
def prune_all_backups(self) -> None:
|
|
"""Prune all backup files for known prefixes on initialization, printing only once."""
|
|
logging.info("Pruning old backups for all known prefixes...")
|
|
prefixes = [
|
|
"sage_cache_info",
|
|
"sage_cache_hash",
|
|
"sage_cache_info-save-error",
|
|
"sage_cache_hash-save-error",
|
|
"sage_cache_info-error",
|
|
"sage_cache_hash-error",
|
|
]
|
|
for prefix in prefixes:
|
|
self.prune_old_backups(prefix)
|
|
|
|
def by_path(self, file_path: str) -> dict:
|
|
"""Get cache info by file path."""
|
|
the_hash = self.hash.get(file_path, "")
|
|
if the_hash:
|
|
return self.info.get(the_hash, {})
|
|
logging.warning(f"No hash found for file: {file_path}")
|
|
return {}
|
|
|
|
def by_hash(self, file_hash: str) -> dict:
|
|
"""Get cache info by file hash."""
|
|
return self.info.get(file_hash, {})
|
|
|
|
def convert_old_cache(self) -> None:
|
|
"""Convert old cache format to new format, splitting into hash and info."""
|
|
logging.info("Converting old cache format to new format.")
|
|
for key, val in self.data.items():
|
|
current_hash = val.get("hash", "")
|
|
if current_hash:
|
|
self.hash[key] = current_hash
|
|
try:
|
|
in_civitai = str_to_bool(val.get("civitai", False))
|
|
except Exception:
|
|
in_civitai = False
|
|
if not in_civitai:
|
|
self.info[current_hash] = val
|
|
for key, val in self.data.items():
|
|
current_hash = val.get("hash", "")
|
|
try:
|
|
in_civitai = str_to_bool(val.get("civitai", False))
|
|
except Exception:
|
|
in_civitai = False
|
|
if current_hash and in_civitai:
|
|
self.info[current_hash] = val
|
|
|
|
def prune_old_backups(self, prefix: str) -> None:
|
|
"""
|
|
Prune old backup files with smart deduplication.
|
|
|
|
Strategy:
|
|
1. Group backups by content hash (exact duplicates)
|
|
2. For each group, keep only the newest file
|
|
3. Group remaining backups by entry count (similar backups)
|
|
4. For similar backups, keep the one with most data
|
|
5. Keep only num_of_backups_to_keep most recent unique backups
|
|
"""
|
|
backups = []
|
|
for f in path_manager.backup_path.iterdir():
|
|
if f.is_file() and f.name.startswith(prefix) and f.suffix == ".json":
|
|
try:
|
|
ctime = f.stat().st_ctime
|
|
file_size = f.stat().st_size
|
|
backups.append((ctime, file_size, f))
|
|
except Exception:
|
|
continue
|
|
|
|
if not backups:
|
|
return
|
|
|
|
backups.sort(reverse=True) # Newest first
|
|
|
|
# Phase 1: Deduplicate exact duplicates (by content hash)
|
|
# Keep newest file from each duplicate group
|
|
hash_to_best = {}
|
|
for ctime, file_size, f in backups:
|
|
try:
|
|
with f.open('rb') as file_obj:
|
|
file_bytes = file_obj.read()
|
|
file_hash = hashlib.sha256(file_bytes).hexdigest()
|
|
|
|
if file_hash not in hash_to_best:
|
|
# First file with this hash, keep it
|
|
hash_to_best[file_hash] = (ctime, file_size, f, file_bytes)
|
|
else:
|
|
# Duplicate found, delete it (we keep the newer one already stored)
|
|
f.unlink(missing_ok=True)
|
|
if f.name in self.backup_manifest:
|
|
del self.backup_manifest[f.name]
|
|
logging.debug(f"Deleted duplicate backup: {f.name}")
|
|
except Exception as e:
|
|
logging.warning(f"Error processing backup {f.name}: {e}")
|
|
continue
|
|
|
|
# Phase 2: Smart similarity deduplication
|
|
# Compare backups by entry count, keep the one with most data
|
|
remaining_backups = []
|
|
for file_hash, (ctime, file_size, f, file_bytes) in hash_to_best.items():
|
|
try:
|
|
data = json.loads(file_bytes.decode('utf-8'))
|
|
entry_count = len(data) if isinstance(data, dict) else 0
|
|
remaining_backups.append((ctime, file_size, entry_count, f, file_hash))
|
|
except Exception as e:
|
|
logging.warning(f"Error parsing backup {f.name}: {e}")
|
|
# Keep it anyway, use file size as proxy for entry count
|
|
remaining_backups.append((ctime, file_size, file_size, f, file_hash))
|
|
|
|
# Sort by entry count (descending) then time (descending)
|
|
remaining_backups.sort(key=lambda x: (x[2], x[0]), reverse=True)
|
|
|
|
# Phase 3: Group by similar entry counts and keep best from each group
|
|
# Two backups are "similar" if their entry counts are within 5% of each other
|
|
similarity_threshold = 0.05
|
|
unique_backups = []
|
|
seen_entry_ranges = []
|
|
|
|
for ctime, file_size, entry_count, f, file_hash in remaining_backups:
|
|
# Check if this backup is similar to any we've already kept
|
|
is_similar = False
|
|
for kept_count in seen_entry_ranges:
|
|
if kept_count == 0 and entry_count == 0:
|
|
is_similar = True
|
|
break
|
|
elif kept_count > 0:
|
|
ratio = abs(entry_count - kept_count) / kept_count
|
|
if ratio <= similarity_threshold:
|
|
is_similar = True
|
|
break
|
|
|
|
if not is_similar:
|
|
# This backup is sufficiently different, keep it
|
|
unique_backups.append((ctime, f))
|
|
seen_entry_ranges.append(entry_count)
|
|
else:
|
|
# Similar to one we already kept, delete it
|
|
try:
|
|
f.unlink(missing_ok=True)
|
|
if f.name in self.backup_manifest:
|
|
del self.backup_manifest[f.name]
|
|
logging.debug(f"Deleted similar backup: {f.name} ({entry_count} entries)")
|
|
except Exception:
|
|
pass
|
|
|
|
# Phase 4: Keep only num_of_backups_to_keep most recent
|
|
unique_backups.sort(reverse=True) # Sort by time, newest first
|
|
keep = unique_backups[:self.num_of_backups_to_keep]
|
|
keep_files = set(f for _, f in keep)
|
|
|
|
for _, f in unique_backups[self.num_of_backups_to_keep:]:
|
|
try:
|
|
f.unlink(missing_ok=True)
|
|
if f.name in self.backup_manifest:
|
|
del self.backup_manifest[f.name]
|
|
logging.debug(f"Deleted old backup (exceeded limit): {f.name}")
|
|
except Exception:
|
|
pass
|
|
|
|
# Save manifest after pruning
|
|
self._save_backup_manifest()
|
|
|
|
def _atomic_write_json(self, path: pathlib.Path, data: Any) -> None:
|
|
"""Write JSON data to a file atomically."""
|
|
file_manager.atomic_write_json(path, data)
|
|
|
|
def _save_json(self, path: pathlib.Path, data: Any, label: str) -> None:
|
|
"""Save data to a JSON file atomically, backing up and pruning old backups on error."""
|
|
try:
|
|
self._atomic_write_json(path, data)
|
|
except Exception as e:
|
|
logging.error(f"Unable to save {label} to {path}: {e}")
|
|
current_date = datetime.datetime.now().strftime("%Y-%m-%dT%H-%M-%S")
|
|
if path.is_file():
|
|
error_prefix = f"{path.stem}-save-error"
|
|
error_backup_path = path_manager.get_backup_file_path(f"{error_prefix}-{current_date}.json")
|
|
try:
|
|
with path.open("r") as src_file, error_backup_path.open("w") as dst_file:
|
|
dst_file.write(src_file.read())
|
|
logging.info(f"Backed up problematic file to {error_backup_path}")
|
|
except Exception as backup_e:
|
|
logging.error(f"Unable to backup error file {path} to {error_backup_path}: {backup_e}")
|
|
|
|
def backup_json(self, backup_prefix: str, data: Any, current_date: str) -> None:
|
|
"""
|
|
Backup data to a JSON file atomically, with smart deduplication.
|
|
|
|
Strategy:
|
|
1. Check manifest first for fast comparison (avoids reading files)
|
|
2. If similar backup exists (within 5% entry count), compare which has more data
|
|
3. Keep the backup with more entries, delete the smaller one
|
|
4. Skip if exact duplicate exists
|
|
"""
|
|
data_json = json.dumps(data, separators=(",", ":"), sort_keys=True, indent=4)
|
|
data_hash = hashlib.sha256(data_json.encode("utf-8")).hexdigest()
|
|
entry_count = len(data) if isinstance(data, dict) else 0
|
|
|
|
# Check manifest first for fast comparison
|
|
similar_backup = None
|
|
similar_entry_count = 0
|
|
|
|
for backup_name, manifest_entry in self.backup_manifest.items():
|
|
if not backup_name.startswith(backup_prefix):
|
|
continue
|
|
|
|
# Check for exact duplicate
|
|
if manifest_entry.get("content_hash") == data_hash:
|
|
logging.debug(f"Skipping backup - exact duplicate exists: {backup_name}")
|
|
return
|
|
|
|
# Check for similar backup (within 5% entry count)
|
|
backup_entry_count = manifest_entry.get("entry_count", 0)
|
|
if entry_count > 0 and backup_entry_count > 0:
|
|
ratio = abs(entry_count - backup_entry_count) / max(entry_count, backup_entry_count)
|
|
if ratio <= 0.05:
|
|
# Similar backup found
|
|
if backup_entry_count > similar_entry_count:
|
|
similar_backup = backup_name
|
|
similar_entry_count = backup_entry_count
|
|
|
|
# If similar backup exists with more data, skip creating new backup
|
|
if similar_backup and similar_entry_count >= entry_count:
|
|
# logging.debug(f"Skipping backup - similar backup with more data exists: {similar_backup} ({similar_entry_count} vs {entry_count} entries)")
|
|
return
|
|
|
|
# If similar backup exists with less data, delete it and create new one
|
|
if similar_backup and similar_entry_count < entry_count:
|
|
similar_path = path_manager.get_backup_file_path(similar_backup)
|
|
try:
|
|
if similar_path.exists():
|
|
similar_path.unlink()
|
|
del self.backup_manifest[similar_backup]
|
|
# logging.info(f"Replaced smaller backup {similar_backup} ({similar_entry_count} entries) with larger backup ({entry_count} entries)")
|
|
except Exception as e:
|
|
logging.warning(f"Failed to delete smaller backup {similar_backup}: {e}")
|
|
|
|
# Create new backup
|
|
safe_date = current_date.replace(":", "-")
|
|
backup_path = path_manager.get_backup_file_path(f"{backup_prefix}-{safe_date}.json")
|
|
try:
|
|
temp_dir = backup_path.parent
|
|
with tempfile.NamedTemporaryFile('w', dir=temp_dir, delete=False, encoding='utf-8') as tf:
|
|
tf.write(data_json)
|
|
tf.flush()
|
|
os.fsync(tf.fileno())
|
|
tempname = tf.name
|
|
os.replace(tempname, backup_path)
|
|
|
|
# Update manifest
|
|
self._update_manifest_for_backup(backup_path, data)
|
|
|
|
logging.debug(f"Created backup: {backup_path.name} ({entry_count} entries)")
|
|
except Exception as e:
|
|
logging.error(f"Unable to backup {backup_prefix} to {backup_path}: {e}")
|
|
|
|
def load_json_file(self, path: pathlib.Path, label: str, current_date: str) -> Optional[Any]:
|
|
"""Load data from a JSON file, backing up the file if an error occurs."""
|
|
try:
|
|
with path.open("r") as read_file:
|
|
data = json.load(read_file)
|
|
return data
|
|
except Exception as e:
|
|
logging.error(f"Unable to load {label} from {path}: {e}")
|
|
if path.is_file():
|
|
safe_date = current_date.replace(":", "-")
|
|
error_prefix = f"{path.stem}-error"
|
|
error_backup_path = path_manager.get_backup_file_path(f"{error_prefix}-{safe_date}.json")
|
|
try:
|
|
with path.open("r") as src_file, error_backup_path.open("w") as dst_file:
|
|
dst_file.write(src_file.read())
|
|
logging.info(f"Backed up problematic file to {error_backup_path}")
|
|
except Exception as backup_e:
|
|
logging.error(f"Unable to backup error file {path} to {error_backup_path}: {backup_e}")
|
|
return None
|
|
|
|
def load(self) -> None:
|
|
"""Load cache from disk only if not already loaded or if file has changed."""
|
|
current_date = datetime.datetime.now().strftime("%Y-%m-%dT%H-%M-%S")
|
|
if not hasattr(self, 'hash_mtime'):
|
|
self.hash_mtime = None
|
|
if not hasattr(self, 'info_mtime'):
|
|
self.info_mtime = None
|
|
if not hasattr(self, 'ollama_mtime'):
|
|
self.ollama_mtime = None
|
|
try:
|
|
hash_needs_reload = False
|
|
info_needs_reload = False
|
|
ollama_needs_reload = False
|
|
|
|
if self.ollama_models_path.is_file():
|
|
ollama_mtime = self.ollama_models_path.stat().st_mtime
|
|
if not self.ollama_models or self.ollama_mtime != ollama_mtime:
|
|
ollama_needs_reload = True
|
|
|
|
if self.hash_path.is_file():
|
|
hash_mtime = self.hash_path.stat().st_mtime
|
|
if not self.hash or self.hash_mtime != hash_mtime:
|
|
hash_needs_reload = True
|
|
if self.info_path.is_file():
|
|
info_mtime = self.info_path.stat().st_mtime
|
|
if not self.info or self.info_mtime != info_mtime:
|
|
info_needs_reload = True
|
|
if self.hash_path.is_file() and self.info_path.is_file():
|
|
if hash_needs_reload:
|
|
#print("Loading hash cache from disk.")
|
|
hash_data = self.load_json_file(self.hash_path, "hash cache", current_date)
|
|
if hash_data is not None:
|
|
self.hash = hash_data
|
|
self.last_hash = copy.deepcopy(self.hash)
|
|
self.hash_mtime = self.hash_path.stat().st_mtime
|
|
self.backup_json("sage_cache_hash", self.hash, current_date)
|
|
else:
|
|
self.hash = {}
|
|
self.last_hash = {}
|
|
self.hash_mtime = None
|
|
if info_needs_reload:
|
|
#print("Loading info cache from disk.")
|
|
info_data = self.load_json_file(self.info_path, "info cache", current_date)
|
|
if info_data is not None:
|
|
self.info = info_data
|
|
self.last_info = copy.deepcopy(self.info)
|
|
self.info_mtime = self.info_path.stat().st_mtime
|
|
self.backup_json("sage_cache_info", self.info, current_date)
|
|
else:
|
|
self.info = {}
|
|
self.last_info = {}
|
|
self.info_mtime = None
|
|
elif self.main_path.is_file():
|
|
data = self.load_json_file(self.main_path, "main cache", current_date)
|
|
if data is not None:
|
|
self.data = data
|
|
self.convert_old_cache()
|
|
else:
|
|
self.data = {}
|
|
if self.ollama_models_path.is_file() and ollama_needs_reload:
|
|
ollama_data = self.load_json_file(self.ollama_models_path, "Ollama models cache", current_date)
|
|
if ollama_data is not None:
|
|
self.ollama_models = ollama_data
|
|
self.last_ollama_models = copy.deepcopy(self.ollama_models)
|
|
self.ollama_mtime = self.ollama_models_path.stat().st_mtime
|
|
self.backup_json("sage_cache_ollama", self.ollama_models, current_date)
|
|
else:
|
|
self.ollama_models = {}
|
|
self.last_ollama_models = {}
|
|
self.ollama_mtime = None
|
|
except Exception as e:
|
|
logging.error(f"Unable to load cache: {e}")
|
|
|
|
def save(self) -> None:
|
|
"""Save cache to disk. Skipped if batch_mode is True."""
|
|
# Skip save if in batch mode
|
|
if self.batch_mode:
|
|
return
|
|
|
|
saved = False
|
|
if self.hash and self.hash != self.last_hash:
|
|
self._save_json(self.hash_path, self.hash, "hash cache")
|
|
self.last_hash = copy.deepcopy(self.hash)
|
|
saved = True
|
|
if self.info and self.info != self.last_info:
|
|
self._save_json(self.info_path, self.info, "info cache")
|
|
self.last_info = copy.deepcopy(self.info)
|
|
saved = True
|
|
if self.ollama_models and self.ollama_models != self.last_ollama_models:
|
|
self._save_json(self.ollama_models_path, self.ollama_models, "Ollama models cache")
|
|
self.last_ollama_models = copy.deepcopy(self.ollama_models)
|
|
saved = True
|
|
if saved:
|
|
self.save_count_since_backup += 1
|
|
logging.info("Saved cache to disk.")
|
|
|
|
def begin_batch(self) -> None:
|
|
"""
|
|
Start a batch operation - suppress saves and backups until end_batch() is called.
|
|
Use this when performing bulk operations to avoid excessive I/O.
|
|
|
|
Example:
|
|
cache.begin_batch()
|
|
try:
|
|
for model in models:
|
|
process_model(model)
|
|
finally:
|
|
cache.end_batch(force_save=True)
|
|
"""
|
|
if self.batch_mode:
|
|
logging.warning("Batch mode already active - ignoring begin_batch() call")
|
|
return
|
|
|
|
self.batch_mode = True
|
|
self.batch_start_changes = self.pending_changes
|
|
logging.info("Batch mode started - saves and backups deferred")
|
|
|
|
def end_batch(self, force_save: bool = True) -> None:
|
|
"""
|
|
End a batch operation - perform a single save and create backup if needed.
|
|
|
|
Args:
|
|
force_save: If True, save even if no changes detected. Default True for safety.
|
|
"""
|
|
if not self.batch_mode:
|
|
logging.warning("Batch mode not active - ignoring end_batch() call")
|
|
return
|
|
|
|
self.batch_mode = False
|
|
changes_in_batch = self.pending_changes - self.batch_start_changes
|
|
|
|
if force_save or changes_in_batch > 0:
|
|
# Perform the deferred save
|
|
saved = False
|
|
if self.hash and self.hash != self.last_hash:
|
|
self._save_json(self.hash_path, self.hash, "hash cache")
|
|
self.last_hash = copy.deepcopy(self.hash)
|
|
saved = True
|
|
if self.info and self.info != self.last_info:
|
|
self._save_json(self.info_path, self.info, "info cache")
|
|
self.last_info = copy.deepcopy(self.info)
|
|
saved = True
|
|
if self.ollama_models and self.ollama_models != self.last_ollama_models:
|
|
self._save_json(self.ollama_models_path, self.ollama_models, "Ollama models cache")
|
|
self.last_ollama_models = copy.deepcopy(self.ollama_models)
|
|
saved = True
|
|
|
|
if saved:
|
|
self.save_count_since_backup += 1
|
|
logging.info(f"Batch save complete ({changes_in_batch} changes)")
|
|
|
|
# Check if backup is needed
|
|
self._create_backups_if_needed()
|
|
else:
|
|
logging.info("Batch mode ended with no changes - no save needed")
|
|
|
|
def _create_backups_if_needed(self) -> None:
|
|
"""
|
|
Create backups based on thresholds (count and time).
|
|
Only creates backup if:
|
|
- save_count_since_backup >= backup_threshold OR
|
|
- time since last backup >= backup_interval_seconds
|
|
"""
|
|
current_time = datetime.datetime.now()
|
|
time_since_backup = (current_time - self.last_backup_time).total_seconds()
|
|
|
|
should_backup = (
|
|
self.save_count_since_backup >= self.backup_threshold or
|
|
time_since_backup >= self.backup_interval_seconds
|
|
)
|
|
|
|
if should_backup:
|
|
current_date = current_time.strftime("%Y-%m-%dT%H-%M-%S")
|
|
|
|
# Create backups
|
|
if self.hash:
|
|
self.backup_json("sage_cache_hash", self.hash, current_date)
|
|
if self.info:
|
|
self.backup_json("sage_cache_info", self.info, current_date)
|
|
if self.ollama_models:
|
|
self.backup_json("sage_cache_ollama", self.ollama_models, current_date)
|
|
|
|
# Reset counters
|
|
self.save_count_since_backup = 0
|
|
self.last_backup_time = current_time
|
|
|
|
logging.info(f"Backups created (saves: {self.save_count_since_backup}, time: {time_since_backup:.0f}s)")
|
|
|
|
def add_entry(self, file_path: str, file_hash: str) -> None:
|
|
self.hash[file_path] = file_hash
|
|
if file_hash not in self.info:
|
|
self.info[file_hash] = {
|
|
"hash": file_hash,
|
|
"lastUsed": "",
|
|
"civitai": "False",
|
|
"filePath": file_path
|
|
}
|
|
self.save()
|
|
|
|
def add_or_update_entry(self, file_path: str, info_dict: dict) -> None:
|
|
"""
|
|
Add or update a cache entry for a given file path.
|
|
Ensures both hash and info are updated together.
|
|
"""
|
|
file_hash = info_dict.get("hash")
|
|
if not file_hash:
|
|
raise ValueError("info_dict must contain a 'hash' key")
|
|
self.hash[file_path] = file_hash
|
|
self.info[file_hash] = info_dict
|
|
|
|
def remove_entry(self, file_path: str) -> None:
|
|
"""
|
|
Remove a cache entry by file path.
|
|
Removes both hash and info if no other file uses the same hash.
|
|
"""
|
|
file_hash = self.hash.get(file_path)
|
|
if file_hash:
|
|
del self.hash[file_path]
|
|
if file_hash not in self.hash.values():
|
|
self.info.pop(file_hash, None)
|
|
|
|
def update_last_used_by_path(self, file_path: str) -> None:
|
|
"""
|
|
Update the 'lastUsed' field for a given file path.
|
|
"""
|
|
file_hash = self.hash.get(file_path)
|
|
if file_hash:
|
|
if file_hash in self.info:
|
|
self.info[file_hash]['lastUsed'] = datetime.datetime.now().isoformat()
|
|
else:
|
|
self.add_entry(file_path, file_hash)
|
|
self.info[file_hash]['lastUsed'] = datetime.datetime.now().isoformat()
|
|
|
|
def update_last_used_by_hash(self, file_hash: str) -> None:
|
|
"""
|
|
Update the 'lastUsed' field for a given file hash.
|
|
"""
|
|
if file_hash in self.info:
|
|
self.info[file_hash]['lastUsed'] = datetime.datetime.now().isoformat()
|
|
|
|
def get_last_used_by_path(self, file_path: str) -> Optional[datetime.datetime]:
|
|
"""
|
|
Get the 'lastUsed' field for a given file path.
|
|
Returns None if the file path is not found.
|
|
"""
|
|
file_hash = self.hash.get(file_path)
|
|
if file_hash and file_hash in self.info:
|
|
last_used_str = self.info[file_hash].get('lastUsed', '')
|
|
if last_used_str:
|
|
return datetime.datetime.fromisoformat(last_used_str)
|
|
return None
|
|
|
|
def get_models_by_model_id(self, model_id: Any) -> List[dict]:
|
|
"""
|
|
Return a list of all model info dicts with the same modelId.
|
|
For each, include a boolean 'latest_version_present' indicating if the update_version_id is present as an id in the cache.
|
|
"""
|
|
id_to_info = {info.get("id"): info for info in self.info.values() if "id" in info}
|
|
models = [info for info in self.info.values() if info.get("modelId") == model_id]
|
|
for info in models:
|
|
update_version_id = info.get("update_version_id")
|
|
info["latest_version_present"] = bool(update_version_id and update_version_id in id_to_info)
|
|
return models
|
|
|
|
# Global cache instance
|
|
cache = SageCache()
|