"""统一生成任务 Celery 任务定义。 使用 BaoDan 的 Celery 实例，通过 autodiscover 或显式导入注册。 任务通过数据库状态机保证幂等，不依赖 Celery 内置重试。 """ import json import logging from datetime import datetime from celery import shared_task logger = logging.getLogger(__name__) def _update_task_status(task_id: str, **kwargs): """更新任务状态（数据库层面）。""" from insurance.db.compat import db from insurance.models.generation_task import GenerationTask task = GenerationTask.query.get(task_id) if not task: return for key, value in kwargs.items(): if hasattr(task, key): setattr(task, key, value) task.heartbeat_at = datetime.now() db.session.commit() return task def _claim_task(task_id: str) -> bool: """尝试领取任务（queued → running），返回是否成功。""" from insurance.db.compat import db from insurance.models.generation_task import GenerationTask result = db.session.execute( db.text("UPDATE insurance_generation_tasks SET status = 'running', " "started_at = NOW(), heartbeat_at = NOW(), " "attempt_count = attempt_count + 1 " "WHERE id = :id AND status = 'queued'"), {"id": task_id}, ) db.session.commit() return result.rowcount > 0 # ─── PPT 解析任务 ────────────────────────────────────────── @shared_task(bind=True, name="insurance.parse_ppt", max_retries=3, default_retry_delay=30) def parse_ppt_task(self, task_id: str): """PPT PDF 解析任务。 幂等：通过数据库状态机保证（queued → running → done/failed）。 """ if not _claim_task(task_id): logger.info(f"任务 {task_id} 已被领取或不在 queued 状态，跳过") return _update_task_status(task_id, stage="extracting", progress=5, message="开始解析 PDF") try: _execute_ppt_parse(task_id) except Exception as exc: logger.error(f"PPT 解析任务失败 [{task_id}]: {exc}", exc_info=True) _update_task_status(task_id, status="failed", error_code="parse_error", error_message=str(exc)[:1000], finished_at=datetime.now()) raise def _execute_ppt_parse(task_id: str): """执行 PPT 解析逻辑。""" import asyncio import os from insurance.db.compat import db from insurance.models.generation_task import GenerationTask from insurance.models.ppt_session import PptSession from insurance.ppt.extraction import ExtractionOrchestrator task = GenerationTask.query.get(task_id) if not task: return session = PptSession.query.get(task.workspace_id) if not session: _update_task_status(task_id, status="failed", error_code="workspace_not_found", error_message="工作区不存在", finished_at=datetime.now()) return files = json.loads(session.files_json) if session.files_json else [] if not files: _update_task_status(task_id, status="failed", error_code="no_files", error_message="没有可解析的 PDF 文件", finished_at=datetime.now()) return orchestrator = ExtractionOrchestrator() extractions = [] total = len(files) for index, file_info in enumerate(files, start=1): filename = file_info.get("name", "") filepath = file_info.get("path", "") plan_type = file_info.get("type", "savings") _update_task_status( task_id, progress=int((index - 1) / total * 80), message=f"正在解析 {filename or f'第 {index} 个文件'}", ) try: result = asyncio.run(orchestrator.extract_plan(filepath, plan_type, force_reparse=True)) extractions.append({ "pdfName": filename, "pdfPath": filepath, "planType": result.plan_type, "status": result.status, "productName": result.product_name, "data": result.data, "error": result.error, "yearCount": len(result.data.get("benefit_illustration", [])) if result.data else 0, }) except Exception as exc: logger.error(f"PDF 解析失败 [{filename}]: {exc}", exc_info=True) extractions.append({ "pdfName": filename, "pdfPath": filepath, "planType": plan_type, "status": "error", "productName": "unknown", "data": None, "error": str(exc), "yearCount": 0, }) # 更新会话 session = PptSession.query.get(task.workspace_id) if not session: return all_failed = all(e.get("status") == "error" for e in extractions) session.extractions_json = json.dumps(extractions, ensure_ascii=False) session.status = "error" if all_failed else "parsed" session.parse_progress = 100 session.parse_finished_at = datetime.now() db.session.commit() # 更新任务状态 _update_task_status( task_id, status="failed" if all_failed else "done", stage="completed", progress=100, message="解析失败" if all_failed else "解析完成", error_code="all_failed" if all_failed else "", error_message="所有文件均处理失败" if all_failed else "", finished_at=datetime.now(), ) # ─── PPT 生成任务 ────────────────────────────────────────── @shared_task(bind=True, name="insurance.generate_ppt", max_retries=3, default_retry_delay=30) def generate_ppt_task(self, task_id: str): """PPT 生成任务。""" if not _claim_task(task_id): logger.info(f"任务 {task_id} 已被领取或不在 queued 状态，跳过") return _update_task_status(task_id, stage="validating", progress=5, message="开始生成 PPT") try: _execute_ppt_generate(task_id) except Exception as exc: logger.error(f"PPT 生成任务失败 [{task_id}]: {exc}", exc_info=True) _update_task_status(task_id, status="failed", error_code="generate_error", error_message=str(exc)[:1000], finished_at=datetime.now()) raise def _execute_ppt_generate(task_id: str): """执行 PPT 生成逻辑。""" import os import uuid as uuid_mod from insurance.db.compat import db from insurance.models.generation_task import GenerationTask from insurance.models.ppt_session import PptSession from insurance.models.ppt_config import PptTemplate, PptCompany from insurance.ppt.normalizer import normalize_savings_plan, normalize_ci_plan, normalize_iul_plan from insurance.ppt.validator import validate_formal_savings_plan, validate_formal_ci_plan, validate_formal_iul_plan from insurance.ppt.renderer import PptRenderer task = GenerationTask.query.get(task_id) if not task: return session = PptSession.query.get(task.workspace_id) if not session: _update_task_status(task_id, status="failed", error_code="workspace_not_found", error_message="工作区不存在", finished_at=datetime.now()) return # 从快照中读取生成参数 snapshot = json.loads(task.input_snapshot_json) if task.input_snapshot_json else {} theme = snapshot.get("theme", "broker") company_id = snapshot.get("companyId", "") use_masked_data = bool(snapshot.get("useMaskedData")) extractions = json.loads(session.extractions_json) if session.extractions_json else [] _update_task_status(task_id, stage="normalizing", progress=10, message="数据归一化中") # 归一化 all_normalized = [] for ext in extractions: if ext.get("status") not in ("success", "partial") or not ext.get("data"): continue ext_data = ext["data"] pdf_path = ext.get("pdfPath") plan_type = (ext.get("planType") or ext_data.get("product_type") or "savings").lower() try: if plan_type == "ci": normalized = normalize_ci_plan(ext_data, pdf_path) issues = validate_formal_ci_plan(normalized) elif plan_type == "iul": normalized = normalize_iul_plan(ext_data, pdf_path) issues = validate_formal_iul_plan(normalized) else: normalized = normalize_savings_plan(ext_data, pdf_path) issues = validate_formal_savings_plan(normalized) errors = [i for i in issues if i.level == "error"] if errors: _update_task_status(task_id, status="failed", error_code="validation_error", error_message=f"数据校验不通过: {'; '.join(e.message for e in errors)}", finished_at=datetime.now()) return all_normalized.append(normalized) except Exception as exc: logger.error(f"归一化失败: {exc}", exc_info=True) if not all_normalized: _update_task_status(task_id, status="failed", error_code="no_valid_data", error_message="无有效提取数据", finished_at=datetime.now()) return _update_task_status(task_id, stage="loading_template", progress=30, message="加载模板") normalized = all_normalized[0] plan_type = normalized.get("kind", "savings") # 加载模板和公司信息 template = PptTemplate.query.filter_by(plan_type=plan_type, style_preset=theme, status=1).first() template_config = template.to_dict() if template else None company_info = None if company_id: company = PptCompany.query.get(company_id) if company: company_info = company.to_dict() _update_task_status(task_id, stage="rendering", progress=50, message="渲染 PPT") # 渲染 renderer = PptRenderer() user_id = task.user_id from insurance.config import get_storage_root output_dir = os.path.join(get_storage_root(), "outputs", "ppt", user_id, task_id) os.makedirs(output_dir, exist_ok=True) output_path = os.path.join(output_dir, f"presentation.pptx") result = renderer.render_enhanced( normalized, output_path, theme=theme, company_id=company_id, company_info=company_info, template_config=template_config, all_products=all_normalized if len(all_normalized) > 1 else None, ) if not result.get("ok"): _update_task_status(task_id, status="failed", error_code="render_error", error_message=f"渲染失败: {result.get('error', '未知错误')}", finished_at=datetime.now()) return _update_task_status(task_id, stage="saving", progress=90, message="保存文件") # 更新会话 session = PptSession.query.get(task.workspace_id) if session: session.ppt_path = output_path session.slide_count = result.get("slideCount", 0) session.status = "done" session.generated_revision = session.draft_revision session.latest_output_path = output_path db.session.commit() # 完成任务 _update_task_status( task_id, status="done", stage="completed", progress=100, message="生成完成", finished_at=datetime.now(), output_json=json.dumps({ "downloadUrl": f"/insurance/ppt/download/{task.workspace_id}", "slideCount": result.get("slideCount", 0), "filePath": output_path, }, ensure_ascii=False), ) # ─── 海报解析任务 ────────────────────────────────────────── @shared_task(bind=True, name="insurance.parse_poster", max_retries=3, default_retry_delay=30) def parse_poster_task(self, task_id: str): """海报计划书解析任务。""" if not _claim_task(task_id): logger.info(f"任务 {task_id} 已被领取或不在 queued 状态，跳过") return _update_task_status(task_id, stage="extracting", progress=5, message="开始解析计划书") try: _execute_poster_parse(task_id) except Exception as exc: logger.error(f"海报解析任务失败 [{task_id}]: {exc}", exc_info=True) _update_task_status(task_id, status="failed", error_code="parse_error", error_message=str(exc)[:1000], finished_at=datetime.now()) raise def _execute_poster_parse(task_id: str): """执行海报解析逻辑。""" import asyncio from insurance.db.compat import db from insurance.models.generation_task import GenerationTask from insurance.models.poster_case_upload import PosterCaseUpload task = GenerationTask.query.get(task_id) if not task: return case_upload_id = task.workspace_id record = PosterCaseUpload.query.get(case_upload_id) if not record: _update_task_status(task_id, status="failed", error_code="workspace_not_found", error_message="记录不存在", finished_at=datetime.now()) return filepath = record.source_file_url if not filepath: _update_task_status(task_id, status="failed", error_code="no_file", error_message="未上传文件", finished_at=datetime.now()) return _update_task_status(task_id, stage="extracting", progress=20, message="解析中") record.parse_status = "parsing" db.session.commit() from insurance.ppt.extraction import ExtractionOrchestrator orchestrator = ExtractionOrchestrator(use_cache=False) parsed = asyncio.run(orchestrator.extract_for_poster(filepath)) record = PosterCaseUpload.query.get(case_upload_id) if record: record.parsed_data = json.dumps(parsed, ensure_ascii=False) record.parse_status = "parsed" db.session.commit() _update_task_status( task_id, status="done", stage="completed", progress=100, message="解析完成", finished_at=datetime.now(), ) # ─── 海报生成任务 ────────────────────────────────────────── @shared_task(bind=True, name="insurance.generate_poster", max_retries=3, default_retry_delay=30) def generate_poster_task(self, task_id: str): """海报图片生成任务。""" if not _claim_task(task_id): logger.info(f"任务 {task_id} 已被领取或不在 queued 状态，跳过") return _update_task_status(task_id, stage="preparing_data", progress=5, message="开始生成海报") try: _execute_poster_generate(task_id) except Exception as exc: logger.error(f"海报生成任务失败 [{task_id}]: {exc}", exc_info=True) _update_task_status(task_id, status="failed", error_code="generate_error", error_message=str(exc)[:1000], finished_at=datetime.now()) raise def _execute_poster_generate(task_id: str): """执行海报生成逻辑。""" import os import re import uuid as uuid_mod from insurance.db.compat import db from insurance.models.generation_task import GenerationTask from insurance.models.poster_record import PosterRecord from insurance.models.poster_template_model import PosterTemplate from insurance.models.ppt_config import PptProduct, PptCompany task = GenerationTask.query.get(task_id) if not task: return record = PosterRecord.query.get(task.workspace_id) if not record: _update_task_status(task_id, status="failed", error_code="workspace_not_found", error_message="工作区不存在", finished_at=datetime.now()) return snapshot = json.loads(task.input_snapshot_json) if task.input_snapshot_json else {} template_id = snapshot.get("templateId") product_id = snapshot.get("productId") copy_content = snapshot.get("copyContent", {}) size = snapshot.get("size", "1024x1792") reference_image = snapshot.get("referenceImage") use_masked_data = bool(snapshot.get("useMaskedData")) _update_task_status(task_id, stage="preparing_data", progress=20, message="准备数据") # 获取模板和产品信息 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) 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 _update_task_status(task_id, stage="building_prompt", progress=40, message="构建 Prompt") # 组装 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, ) _update_task_status(task_id, stage="requesting_image", progress=60, message="请求图片生成") # 生成图片 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" _update_task_status(task_id, stage="saving", progress=85, message="保存文件") # 保存文件 from insurance.config import get_storage_root output_dir = os.path.join(get_storage_root(), "outputs", "posters", task_id) os.makedirs(output_dir, exist_ok=True) filename = f"poster.png" filepath = os.path.join(output_dir, filename) with open(filepath, "wb") as f: f.write(image_bytes) # 更新记录 record = PosterRecord.query.get(task.workspace_id) if record: 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() record.generated_revision = record.draft_revision db.session.commit() # 完成任务 _update_task_status( task_id, status="done", stage="completed", progress=100, message="生成完成", finished_at=datetime.now(), output_json=json.dumps({ "downloadUrl": f"/insurance/poster/download/{task.workspace_id}", "filePath": filepath, }, ensure_ascii=False), )