Spaces:
Running
Running
File size: 6,002 Bytes
e6404d0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 | import time
import uuid
import json
import os
_HISTORY_FILE = os.path.join(os.path.dirname(os.path.dirname(__file__)), "data", "session_history.json")
class HistoryEntry:
__slots__ = ("id", "timestamp", "prompt", "model", "rating", "num_variations",
"creativity", "weight_mode", "results", "neg_results",
"categories", "seed", "liked_indices")
def __init__(self, prompt: str, results: list[str], model: str = "anima",
rating: str = "pg", num_variations: int = 5,
creativity: str = "medium", weight_mode: str = "off",
neg_results: list[str] | None = None,
categories: list[str] | None = None,
seed: int | None = None):
self.id = uuid.uuid4().hex[:8]
self.timestamp = time.time()
self.prompt = prompt
self.model = model
self.rating = rating
self.num_variations = num_variations
self.creativity = creativity
self.weight_mode = weight_mode
self.results = list(results)
self.neg_results = list(neg_results) if neg_results else []
self.categories = list(categories) if categories else []
self.seed = seed
self.liked_indices: list[int] = []
@classmethod
def from_dict(cls, d: dict) -> "HistoryEntry":
entry = cls.__new__(cls)
entry.id = d.get("id", uuid.uuid4().hex[:8])
entry.timestamp = d.get("timestamp", time.time())
entry.prompt = d.get("prompt", "")
entry.model = d.get("model", "anima")
entry.rating = d.get("rating", "pg")
entry.num_variations = d.get("num_variations", 5)
entry.creativity = d.get("creativity", "medium")
entry.weight_mode = d.get("weight_mode", "off")
entry.results = list(d.get("results", []))
entry.neg_results = list(d.get("neg_results", []))
entry.categories = list(d.get("categories", []))
entry.seed = d.get("seed")
entry.liked_indices = list(d.get("liked_indices", []))
return entry
def toggle_like(self, index: int) -> bool:
if index in self.liked_indices:
self.liked_indices.remove(index)
return False
self.liked_indices.append(index)
return True
def is_liked(self, index: int) -> bool:
return index in self.liked_indices
def get_liked_results(self) -> list[tuple[int, str]]:
return [(i, self.results[i]) for i in self.liked_indices if i < len(self.results)]
def to_dict(self) -> dict:
return {
"id": self.id,
"timestamp": self.timestamp,
"prompt": self.prompt,
"model": self.model,
"rating": self.rating,
"num_variations": self.num_variations,
"creativity": self.creativity,
"weight_mode": self.weight_mode,
"results": self.results,
"neg_results": self.neg_results,
"categories": self.categories,
"seed": self.seed,
"liked_indices": self.liked_indices,
}
class SessionHistory:
def __init__(self, max_entries: int = 50, file_path: str | None = None):
self._entries: list[HistoryEntry] = []
self._max_entries = max_entries
self._file_path = file_path or _HISTORY_FILE
self._load()
def persist(self):
self._save()
def toggle_like(self, entry_index: int, result_index: int) -> bool:
if entry_index < 0 or entry_index >= len(self._entries):
return False
liked = self._entries[entry_index].toggle_like(result_index)
self._save()
return liked
def add(self, entry: HistoryEntry) -> str:
self._entries.insert(0, entry)
if len(self._entries) > self._max_entries:
self._entries.pop()
self._save()
return entry.id
def get(self, entry_id: str) -> HistoryEntry | None:
for e in self._entries:
if e.id == entry_id:
return e
return None
def get_all(self) -> list[HistoryEntry]:
return list(self._entries)
def get_favorites(self) -> list[tuple[HistoryEntry, int, str]]:
items = []
for entry in self._entries:
for idx, text in entry.get_liked_results():
items.append((entry, idx, text))
return items
def clear(self):
self._entries.clear()
self._save()
def _save(self):
data = [e.to_dict() for e in self._entries]
tmp_path = self._file_path + ".tmp"
try:
with open(tmp_path, "w", encoding="utf-8") as f:
json.dump(data, f, ensure_ascii=False, indent=2)
if os.path.exists(self._file_path):
os.replace(self._file_path, self._file_path + ".bak")
os.replace(tmp_path, self._file_path)
except (OSError, TypeError):
try:
if os.path.exists(tmp_path):
os.remove(tmp_path)
except OSError:
pass
def _load(self):
# Try primary, then .bak (recovers from an interrupted write), then reset.
for path in (self._file_path, self._file_path + ".bak"):
if not os.path.exists(path):
continue
try:
with open(path, "r", encoding="utf-8") as f:
data = json.load(f)
self._entries = [HistoryEntry.from_dict(d) for d in data[:self._max_entries]]
return
except (json.JSONDecodeError, UnicodeDecodeError, ValueError, OSError):
continue
self._entries = []
def __len__(self) -> int:
return len(self._entries)
_history = SessionHistory()
def get_history() -> SessionHistory:
return _history
def add_to_history(prompt: str, results: list[str], **kwargs) -> str:
entry = HistoryEntry(prompt, results, **kwargs)
return _history.add(entry)
|