195 lines
6.6 KiB
Python
195 lines
6.6 KiB
Python
import os
|
|
import json
|
|
import datetime
|
|
import time
|
|
import config
|
|
import threading
|
|
|
|
SAFETY_SETTINGS = [
|
|
{"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_NONE"},
|
|
{"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_NONE"},
|
|
{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "BLOCK_NONE"},
|
|
{"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_NONE"},
|
|
]
|
|
|
|
# Thread-local storage for logging context
|
|
_log_context = threading.local()
|
|
|
|
def set_log_file(filepath):
|
|
_log_context.log_file = filepath
|
|
|
|
def set_log_callback(callback):
|
|
_log_context.callback = callback
|
|
|
|
def set_progress_callback(callback):
|
|
_log_context.progress_callback = callback
|
|
|
|
def update_progress(percent):
|
|
if getattr(_log_context, 'progress_callback', None):
|
|
try: _log_context.progress_callback(percent)
|
|
except: pass
|
|
|
|
def clean_json(text):
|
|
text = text.replace("```json", "").replace("```", "").strip()
|
|
# Robust extraction: find first { or [ and last } or ]
|
|
start_obj = text.find('{')
|
|
start_arr = text.find('[')
|
|
if start_obj == -1 and start_arr == -1: return text
|
|
if start_obj != -1 and (start_arr == -1 or start_obj < start_arr):
|
|
return text[start_obj:text.rfind('}')+1]
|
|
else:
|
|
return text[start_arr:text.rfind(']')+1]
|
|
|
|
def sanitize_filename(name):
|
|
"""Sanitizes a string to be safe for filenames."""
|
|
if not name: return "Untitled"
|
|
safe = "".join([c for c in name if c.isalnum() or c=='_']).replace(" ", "_")
|
|
return safe if safe else "Untitled"
|
|
|
|
def chapter_sort_key(ch):
|
|
"""Sort key for chapters handling integers, strings, Prologue, and Epilogue."""
|
|
num = ch.get('num', 0)
|
|
if isinstance(num, int): return num
|
|
if isinstance(num, str) and num.isdigit(): return int(num)
|
|
s = str(num).lower().strip()
|
|
if 'prologue' in s: return -1
|
|
if 'epilogue' in s: return 9999
|
|
return 999
|
|
|
|
def get_sorted_book_folders(run_dir):
|
|
"""Returns a list of book folder names in a run directory, sorted numerically."""
|
|
if not os.path.exists(run_dir): return []
|
|
subdirs = [d for d in os.listdir(run_dir) if os.path.isdir(os.path.join(run_dir, d)) and d.startswith("Book_")]
|
|
def sort_key(d):
|
|
parts = d.split('_')
|
|
if len(parts) > 1 and parts[1].isdigit(): return int(parts[1])
|
|
return 0
|
|
return sorted(subdirs, key=sort_key)
|
|
|
|
# --- SHARED UTILS ---
|
|
def log(phase, msg):
|
|
timestamp = datetime.datetime.now().strftime('%H:%M:%S')
|
|
line = f"[{timestamp}] {phase:<15} | {msg}"
|
|
print(line)
|
|
|
|
# Write to thread-specific log file if set
|
|
if getattr(_log_context, 'log_file', None):
|
|
with open(_log_context.log_file, "a", encoding="utf-8") as f:
|
|
f.write(line + "\n")
|
|
|
|
# Trigger callback if set (e.g. for Database logging)
|
|
if getattr(_log_context, 'callback', None):
|
|
try: _log_context.callback(phase, msg)
|
|
except: pass
|
|
|
|
def load_json(path):
|
|
return json.load(open(path, 'r')) if os.path.exists(path) else None
|
|
|
|
def create_default_personas():
|
|
# Initialize empty personas file if it doesn't exist
|
|
if not os.path.exists(config.PERSONAS_DIR): os.makedirs(config.PERSONAS_DIR)
|
|
if not os.path.exists(config.PERSONAS_FILE):
|
|
try:
|
|
with open(config.PERSONAS_FILE, 'w') as f: json.dump({}, f, indent=2)
|
|
except: pass
|
|
|
|
def get_length_presets():
|
|
"""Returns a dict mapping Label -> Settings for use in main.py"""
|
|
presets = {}
|
|
for k, v in config.LENGTH_DEFINITIONS.items():
|
|
presets[v['label']] = v
|
|
return presets
|
|
|
|
def log_image_attempt(folder, img_type, prompt, filename, status, error=None, score=None, critique=None):
|
|
log_path = os.path.join(folder, "image_log.json")
|
|
entry = {
|
|
"timestamp": int(time.time()),
|
|
"type": img_type,
|
|
"prompt": prompt,
|
|
"filename": filename,
|
|
"status": status,
|
|
"error": str(error) if error else None,
|
|
"score": score,
|
|
"critique": critique
|
|
}
|
|
data = []
|
|
if os.path.exists(log_path):
|
|
try:
|
|
with open(log_path, 'r') as f: data = json.load(f)
|
|
except:
|
|
pass
|
|
data.append(entry)
|
|
with open(log_path, 'w') as f: json.dump(data, f, indent=2)
|
|
|
|
def get_run_folder(base_name):
|
|
if not os.path.exists(base_name): os.makedirs(base_name)
|
|
runs = [d for d in os.listdir(base_name) if d.startswith("run_")]
|
|
next_num = max([int(r.split("_")[1]) for r in runs if r.split("_")[1].isdigit()] + [0]) + 1
|
|
folder = os.path.join(base_name, f"run_{next_num}")
|
|
os.makedirs(folder)
|
|
return folder
|
|
|
|
def get_latest_run_folder(base_name):
|
|
if not os.path.exists(base_name): return None
|
|
runs = [d for d in os.listdir(base_name) if d.startswith("run_")]
|
|
if not runs: return None
|
|
runs.sort(key=lambda x: int(x.split('_')[1]) if x.split('_')[1].isdigit() else 0)
|
|
return os.path.join(base_name, runs[-1])
|
|
|
|
def log_usage(folder, model_label, usage_metadata=None, image_count=0):
|
|
if not folder or not os.path.exists(folder): return
|
|
|
|
log_path = os.path.join(folder, "usage_log.json")
|
|
|
|
entry = {
|
|
"timestamp": int(time.time()),
|
|
"model": model_label,
|
|
"input_tokens": 0,
|
|
"output_tokens": 0,
|
|
"images": image_count
|
|
}
|
|
|
|
if usage_metadata:
|
|
try:
|
|
entry["input_tokens"] = usage_metadata.prompt_token_count
|
|
entry["output_tokens"] = usage_metadata.candidates_token_count
|
|
except: pass
|
|
|
|
data = {"log": [], "totals": {"input_tokens": 0, "output_tokens": 0, "images": 0, "est_cost_usd": 0.0}}
|
|
|
|
if os.path.exists(log_path):
|
|
try:
|
|
loaded = json.load(open(log_path, 'r'))
|
|
if isinstance(loaded, list): data["log"] = loaded
|
|
else: data = loaded
|
|
except: pass
|
|
|
|
data["log"].append(entry)
|
|
|
|
# Recalculate totals
|
|
t_in = sum(x.get('input_tokens', 0) for x in data["log"])
|
|
t_out = sum(x.get('output_tokens', 0) for x in data["log"])
|
|
t_img = sum(x.get('images', 0) for x in data["log"])
|
|
|
|
cost = 0.0
|
|
for x in data["log"]:
|
|
m = x.get('model', '').lower()
|
|
i = x.get('input_tokens', 0)
|
|
o = x.get('output_tokens', 0)
|
|
imgs = x.get('images', 0)
|
|
|
|
if 'flash' in m:
|
|
cost += (i / 1_000_000 * 0.075) + (o / 1_000_000 * 0.30)
|
|
elif 'pro' in m or 'logic' in m:
|
|
cost += (i / 1_000_000 * 3.50) + (o / 1_000_000 * 10.50)
|
|
elif 'imagen' in m or imgs > 0:
|
|
cost += (imgs * 0.04)
|
|
|
|
data["totals"] = {
|
|
"input_tokens": t_in,
|
|
"output_tokens": t_out,
|
|
"images": t_img,
|
|
"est_cost_usd": round(cost, 4)
|
|
}
|
|
|
|
with open(log_path, 'w') as f: json.dump(data, f, indent=2) |