"""海报后台任务管理。 使用与 PPT parse_worker 相同的后台线程 + Redis 锁模式。 任务状态通过数据库字段追踪，前端轮询获取进度。 """ import json import logging import re import threading import uuid import os from datetime import datetime from insurance.db.compat import db logger = logging.getLogger(__name__) _local_locks: set = set() _local_locks_guard = threading.Lock() _redis_locks: set = set() # 超过此时间仍为 queued/generating 的任务视为过期（秒） STALE_TASK_TIMEOUT = 600 # 10 分钟 def recover_stale_tasks(): """启动时恢复过期任务：将长时间 queued/generating 的任务标记为 failed。 应在应用启动时调用。 """ from insurance.models.poster_record import PosterRecord from insurance.models.poster_case_upload import PosterCaseUpload cutoff = datetime.now().timestamp() - STALE_TASK_TIMEOUT # 恢复过期的海报生成任务 stale_records = PosterRecord.query.filter( PosterRecord.task_status.in_(["queued", "generating"]), PosterRecord.created_at < datetime.fromtimestamp(cutoff), ).all() for record in stale_records: record.task_status = "failed" record.task_error = "任务因服务重启而中断，请重新生成" record.finished_at = datetime.now() logger.warning(f"恢复过期海报任务: record_id={record.id}") # 恢复过期的计划书解析任务 stale_cases = PosterCaseUpload.query.filter( PosterCaseUpload.parse_status.in_(["queued", "parsing"]), PosterCaseUpload.created_at < datetime.fromtimestamp(cutoff), ).all() for case in stale_cases: case.parse_status = "failed" logger.warning(f"恢复过期解析任务: case_id={case.id}") if stale_records or stale_cases: db.session.commit() logger.info(f"已恢复 {len(stale_records)} 个海报任务, {len(stale_cases)} 个解析任务") # ─── 计划书解析任务 ─────────────────────────────────────── def start_case_parse_task(app, case_upload_id: int) -> bool: """启动后台计划书解析任务，返回是否新启动。""" lock_key = f"poster_case_parse:{case_upload_id}" if not _acquire_lock(lock_key): return False thread = threading.Thread( target=_run_case_parse_task, args=(app, case_upload_id, lock_key), daemon=True, ) thread.start() return True def _run_case_parse_task(app, case_upload_id: int, lock_key: str): with app.app_context(): try: _execute_case_parse(case_upload_id) except Exception as exc: logger.error(f"计划书解析任务失败 [{case_upload_id}]: {exc}", exc_info=True) _mark_case_failed(case_upload_id, str(exc)) finally: _release_lock(lock_key) db.session.remove() def _execute_case_parse(case_upload_id: int): from insurance.models.poster_case_upload import PosterCaseUpload record = PosterCaseUpload.query.get(case_upload_id) if not record: return filepath = record.source_file_url if not filepath or not os.path.exists(filepath): record.parse_status = "failed" db.session.commit() return record.parse_status = "parsing" db.session.commit() from insurance.ppt.extraction import ExtractionOrchestrator orchestrator = ExtractionOrchestrator(use_cache=False) import asyncio parsed = asyncio.run(orchestrator.extract_for_poster(filepath)) record = PosterCaseUpload.query.get(case_upload_id) if not record: return record.parsed_data = json.dumps(parsed, ensure_ascii=False) record.parse_status = "parsed" db.session.commit() def _mark_case_failed(case_upload_id: int, error: str): from insurance.models.poster_case_upload import PosterCaseUpload record = PosterCaseUpload.query.get(case_upload_id) if not record: return record.parse_status = "failed" db.session.commit() # ─── 海报图片生成任务 ────────────────────────────────────── def start_poster_generate_task(app, record_id: int, user_id: str, data: dict) -> bool: """启动后台海报图片生成任务，返回是否新启动。""" lock_key = f"poster_generate:{record_id}" if not _acquire_lock(lock_key): return False thread = threading.Thread( target=_run_poster_generate_task, args=(app, record_id, user_id, data, lock_key), daemon=True, ) thread.start() return True def _run_poster_generate_task(app, record_id: int, user_id: str, data: dict, lock_key: str): with app.app_context(): try: _execute_poster_generate(record_id, user_id, data) except Exception as exc: logger.error(f"海报生成任务失败 [{record_id}]: {exc}", exc_info=True) _mark_poster_failed(record_id, str(exc)) finally: _release_lock(lock_key) db.session.remove() def _execute_poster_generate(record_id: int, user_id: str, data: dict): from insurance.models.poster_record import PosterRecord from insurance.models.poster_template_model import PosterTemplate from insurance.models.ppt_config import PptProduct, PptCompany record = PosterRecord.query.get(record_id) if not record: return record.task_status = "generating" record.task_progress = 10 record.started_at = datetime.now() db.session.commit() template_id = data.get("templateId") product_id = data.get("productId") copy_content = data.get("copyContent", {}) size = data.get("size", "1024x1792") reference_image = data.get("referenceImage") use_masked_data = bool(data.get("useMaskedData")) # 获取模板和产品信息 poster_template = PosterTemplate.query.get(template_id) if template_id else None product = PptProduct.query.get(product_id) if product_id else None company = PptCompany.query.get(product.company_id) if product else None # 脱敏处理 if use_masked_data: from insurance.ppt.masking import apply_product_mask, apply_company_mask, mask_text product_dict = product.to_dict() if product else None company_dict = company.to_dict() if company else None if product_dict: apply_product_mask(product_dict, True) if company_dict: apply_company_mask(company_dict, True) # 替换 copy_content 中的产品名和保司名 if product and company: real_name = product.display_name masked_name = (product_dict or {}).get("displayName", "") real_company = company.display_name masked_company = (company_dict or {}).get("displayName", "") replacements = {} if real_name and masked_name and real_name != masked_name: replacements[real_name] = masked_name if real_company and masked_company and real_company != masked_company: replacements[real_company] = masked_company if replacements: for key in copy_content: if isinstance(copy_content[key], str): copy_content[key] = mask_text(copy_content[key], replacements) else: product_dict = product.to_dict() if product else None company_dict = {"displayName": company.display_name} if company else None record.task_progress = 30 db.session.commit() # 组装 prompt from insurance.poster.image_generator import PosterImageGenerator generator = PosterImageGenerator() prompt = generator.build_prompt( template=poster_template.to_dict() if poster_template else None, product=product_dict, company=company_dict, copy=copy_content, size=size, ) record.task_progress = 50 db.session.commit() # 生成图片 generation_mode = "ai" provider_info = {} try: image_bytes, provider_info = generator.generate(prompt, size=size, reference_image=reference_image) except Exception as e: logger.warning(f"图片 API 失败，使用降级方案: {e}") from insurance.poster.image_generator import generate_fallback image_bytes = generate_fallback(copy_content, size=size) generation_mode = "fallback" record.task_progress = 80 db.session.commit() # 保存文件（使用持久化存储） from insurance.config import get_storage_root output_dir = os.path.join(get_storage_root(), "outputs", "posters") os.makedirs(output_dir, exist_ok=True) safe_uid = re.sub(r'[/\\.]', '_', user_id) filename = f"poster_{safe_uid}_{uuid.uuid4().hex[:8]}.png" filepath = os.path.join(output_dir, filename) with open(filepath, "wb") as f: f.write(image_bytes) # 更新记录 record = PosterRecord.query.get(record_id) if not record: return record.export_url = filepath record.export_format = "png" record.generation_mode = generation_mode record.image_provider = provider_info.get("provider", "") record.image_model = provider_info.get("model", "") record.prompt_used = prompt[:2000] if prompt else None record.task_status = "done" record.task_progress = 100 record.finished_at = datetime.now() db.session.commit() def _mark_poster_failed(record_id: int, error: str): from insurance.models.poster_record import PosterRecord record = PosterRecord.query.get(record_id) if not record: return record.task_status = "failed" record.task_error = error[:1000] record.task_progress = 100 record.finished_at = datetime.now() db.session.commit() # ─── 锁管理 ──────────────────────────────────────────────── def _acquire_lock(key: str) -> bool: """获取任务锁（优先 Redis，降级本地内存）。""" redis_key = f"poster_task_lock:{key}" try: from insurance.db.compat import redis_client if redis_client and redis_client.set(redis_key, "1", nx=True, ex=1800): _redis_locks.add(key) return True if redis_client: return False except Exception: pass with _local_locks_guard: if key in _local_locks: return False _local_locks.add(key) return True def _release_lock(key: str): """释放任务锁。""" if key in _redis_locks: try: from insurance.db.compat import redis_client if redis_client: redis_client.delete(f"poster_task_lock:{key}") except Exception: pass _redis_locks.discard(key) with _local_locks_guard: _local_locks.discard(key)