import json import os import logging from flask import current_app from sqlalchemy.orm.attributes import flag_modified from models import db, Character, Outfit, Action, Style, Scene, Detailer, Look, Checkpoint from services.llm import load_prompt, call_llm from services.sync import _sync_nsfw_from_tags from services.job_queue import _enqueue_task logger = logging.getLogger('gaze') # Map category string to (model class, id_field, config_dir_key) _CATEGORY_MAP = { 'characters': (Character, 'character_id', 'CHARACTERS_DIR'), 'outfits': (Outfit, 'outfit_id', 'CLOTHING_DIR'), 'actions': (Action, 'action_id', 'ACTIONS_DIR'), 'styles': (Style, 'style_id', 'STYLES_DIR'), 'scenes': (Scene, 'scene_id', 'SCENES_DIR'), 'detailers': (Detailer, 'detailer_id', 'DETAILERS_DIR'), 'looks': (Look, 'look_id', 'LOOKS_DIR'), } # Fields to preserve from the original data (never overwritten by LLM output) _PRESERVE_KEYS = { 'lora', 'participants', 'suppress_wardrobe', 'character_id', 'character_name', 'outfit_id', 'outfit_name', 'action_id', 'action_name', 'style_id', 'style_name', 'scene_id', 'scene_name', 'detailer_id', 'detailer_name', 'look_id', 'look_name', } def register_routes(app): @app.route('/api///regenerate_tags', methods=['POST']) def regenerate_tags(category, slug): if category not in _CATEGORY_MAP: return {'error': f'Unknown category: {category}'}, 400 model_class, id_field, dir_key = _CATEGORY_MAP[category] entity = model_class.query.filter_by(slug=slug).first() if not entity: return {'error': 'Not found'}, 404 system_prompt = load_prompt('regenerate_tags_system.txt') if not system_prompt: return {'error': 'Regenerate tags system prompt not found'}, 500 original_data = entity.data.copy() try: prompt = ( f"Regenerate the prompt tags for this {category.rstrip('s')} resource.\n" f"Current JSON:\n{json.dumps(original_data, indent=2)}" ) llm_response = call_llm(prompt, system_prompt) clean_json = llm_response.replace('```json', '').replace('```', '').strip() new_data = json.loads(clean_json) except Exception as e: logger.exception("Regenerate tags LLM error for %s/%s", category, slug) return {'error': f'LLM error: {str(e)}'}, 500 # Preserve protected fields from original for key in _PRESERVE_KEYS: if key in original_data: new_data[key] = original_data[key] # Update DB entity.data = new_data flag_modified(entity, 'data') _sync_nsfw_from_tags(entity, new_data) db.session.commit() # Write back to JSON file if entity.filename: file_path = os.path.join(current_app.config[dir_key], entity.filename) with open(file_path, 'w') as f: json.dump(new_data, f, indent=2) return {'success': True, 'data': new_data} @app.route('/admin/migrate_tags', methods=['POST']) def migrate_tags(): """One-time migration: convert old list-format tags to new dict format.""" migrated = 0 for category, (model_class, id_field, dir_key) in _CATEGORY_MAP.items(): entities = model_class.query.all() for entity in entities: tags = entity.data.get('tags') if isinstance(tags, list) or tags is None: new_data = entity.data.copy() new_data['tags'] = {'nsfw': False} entity.data = new_data flag_modified(entity, 'data') # Write back to JSON file if entity.filename: file_path = os.path.join(current_app.config[dir_key], entity.filename) try: with open(file_path, 'w') as f: json.dump(new_data, f, indent=2) except Exception as e: logger.warning("Could not write %s: %s", file_path, e) migrated += 1 # Also handle checkpoints for ckpt in Checkpoint.query.all(): data = ckpt.data or {} tags = data.get('tags') if isinstance(tags, list) or tags is None: new_data = data.copy() new_data['tags'] = {'nsfw': False} ckpt.data = new_data flag_modified(ckpt, 'data') migrated += 1 db.session.commit() logger.info("Migrated %d resources from list tags to dict tags", migrated) return {'success': True, 'migrated': migrated} def _make_regen_task(category, slug, name, system_prompt): """Factory: create a tag regeneration task function for one entity.""" def task_fn(job): model_class, id_field, dir_key = _CATEGORY_MAP[category] entity = model_class.query.filter_by(slug=slug).first() if not entity: raise Exception(f'{category}/{slug} not found') original_data = entity.data.copy() prompt = ( f"Regenerate the prompt tags for this {category.rstrip('s')} resource.\n" f"Current JSON:\n{json.dumps(original_data, indent=2)}" ) llm_response = call_llm(prompt, system_prompt) clean_json = llm_response.replace('```json', '').replace('```', '').strip() new_data = json.loads(clean_json) for key in _PRESERVE_KEYS: if key in original_data: new_data[key] = original_data[key] entity.data = new_data flag_modified(entity, 'data') _sync_nsfw_from_tags(entity, new_data) db.session.commit() if entity.filename: file_path = os.path.join(current_app.config[dir_key], entity.filename) with open(file_path, 'w') as f: json.dump(new_data, f, indent=2) job['result'] = {'entity': name, 'status': 'updated'} return task_fn @app.route('/admin/bulk_regenerate_tags/', methods=['POST']) def bulk_regenerate_tags_category(category): """Queue LLM tag regeneration for all resources in a single category.""" if category not in _CATEGORY_MAP: return {'error': f'Unknown category: {category}'}, 400 system_prompt = load_prompt('regenerate_tags_system.txt') if not system_prompt: return {'error': 'Regenerate tags system prompt not found'}, 500 model_class, id_field, dir_key = _CATEGORY_MAP[category] entities = model_class.query.all() job_ids = [] for entity in entities: job = _enqueue_task( f"Regen tags: {entity.name} ({category})", _make_regen_task(category, entity.slug, entity.name, system_prompt) ) job_ids.append(job['id']) return {'success': True, 'queued': len(job_ids), 'job_ids': job_ids} @app.route('/admin/bulk_regenerate_tags', methods=['POST']) def bulk_regenerate_tags(): """Queue LLM tag regeneration for all resources across all categories.""" system_prompt = load_prompt('regenerate_tags_system.txt') if not system_prompt: return {'error': 'Regenerate tags system prompt not found'}, 500 job_ids = [] for category, (model_class, id_field, dir_key) in _CATEGORY_MAP.items(): entities = model_class.query.all() for entity in entities: job = _enqueue_task( f"Regen tags: {entity.name} ({category})", _make_regen_task(category, entity.slug, entity.name, system_prompt) ) job_ids.append(job['id']) return {'success': True, 'queued': len(job_ids), 'job_ids': job_ids}