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)