import asyncio import json import time from datetime import datetime from sqlalchemy import select from database import async_session from models.ai_model import AiModel from models.document import Document from models.generation_log import GenerationLog from models.paragraph import Paragraph from models.template import Template from models.template_block import TemplateBlock from services.ai_service import call_ai from services.file_summary import summarize_minio_files generation_progress: dict[int, dict] = {} generation_cancel_flags: dict[int, bool] = {} def build_mock_content(paragraph: Paragraph) -> dict: if paragraph.output_format == "table": return { "content": [ { "type": "table", "title": paragraph.title, "headers": ["字段", "内容"], "rows": [ ["段落标题", paragraph.title], ["生成说明", paragraph.prompt_text or "根据模板内容生成"], ], } ] } blocks = [ { "type": "text", "text": f"这是“{paragraph.title}”的示例生成内容，可用于前端联调与流程验证。" } ] if paragraph.content: blocks.append({"type": "text", "text": f"模板上下文：{paragraph.content[:200]}"}) if paragraph.need_prompt and paragraph.prompt_text: blocks.append({"type": "text", "text": f"预设提示词：{paragraph.prompt_text[:200]}"}) return {"content": blocks} async def get_effective_model(paragraph: Paragraph) -> AiModel | None: async with async_session() as db: if paragraph.model_id: model = await db.get(AiModel, paragraph.model_id) if model is not None and model.status == "enabled": return model result = await db.execute( select(AiModel).where(AiModel.status == "enabled").order_by(AiModel.id.asc()).limit(1) ) return result.scalars().first() def update_progress(document_id: int, **kwargs): state = generation_progress.setdefault( document_id, {"percent": 0, "status": "pending", "message": "等待中", "done": 0, "total": 0}, ) state.update(kwargs) def request_cancel(document_id: int): generation_cancel_flags[document_id] = True update_progress(document_id, status="cancelling", message="正在取消...") def is_cancel_requested(document_id: int) -> bool: return generation_cancel_flags.get(document_id, False) async def run_generation(document_id: int, template_id: int): async with async_session() as db: document = await db.get(Document, document_id) template = await db.get(Template, template_id) if document is None or template is None: update_progress(document_id, status="failed", message="生成任务初始化失败") return result = await db.execute( select(Paragraph) .where(Paragraph.template_id == template_id) .order_by(Paragraph.sort_index.asc(), Paragraph.id.asc()) ) paragraphs = result.scalars().all() total = len(paragraphs) request_payload = {} if document.request_payload_json: try: request_payload = json.loads(document.request_payload_json) except Exception: request_payload = {} file_map = request_payload.get("file_map", {}) if isinstance(request_payload, dict) else {} update_progress(document_id, status="generating", total=total, done=0, percent=0, message="开始生成...") done_count = 0 failed_count = 0 try: for index, paragraph in enumerate(paragraphs, start=1): if is_cancel_requested(document_id): document.status = "cancelled" document.error = "用户已取消生成" await db.commit() update_progress(document_id, status="cancelled", percent=min(99, int(done_count / max(total, 1) * 100)), message="已取消生成", done=done_count) return if paragraph.edit_mode == "manual": content = {"content": [{"type": "text", "text": paragraph.content or "该段落为人工编辑模式。"}]} status = "success" duration = 0 error_message = "" model_id = paragraph.model_id else: start = time.perf_counter() model = await get_effective_model(paragraph) model_id = model.id if model is not None else paragraph.model_id try: selected_file_paths = file_map.get(str(paragraph.id), []) file_summaries = [] if selected_file_paths: file_name_mapping = {} for paragraph_item in request_payload.get("paragraphs", []): for selected_file in paragraph_item.get("selected_files", []): file_path = selected_file.get("file_path") if file_path in selected_file_paths: file_name_mapping[file_path] = selected_file.get("file_name") file_summaries = await asyncio.to_thread( summarize_minio_files, selected_file_paths, file_name_mapping, ) if model is None: content = build_mock_content(paragraph) else: setattr(paragraph, "enable_reasoning", bool(model.enable_reasoning)) result_data = await call_ai(paragraph, model, file_summaries) content = result_data.content status = "success" error_message = "" except Exception as error: content = build_mock_content(paragraph) status = "failed" error_message = str(error) failed_count += 1 duration = round(time.perf_counter() - start, 4) log = GenerationLog( document_id=document.id, paragraph_id=paragraph.id, model_id=model_id, status=status, content=json.dumps(content, ensure_ascii=False), duration=duration, error_msg=error_message, ) db.add(log) done_count += 1 document.para_count_done = done_count percent = int(done_count / max(total, 1) * 100) update_progress( document_id, status="generating", percent=percent, done=done_count, total=total, current_paragraph=paragraph.title, message=f"正在生成：{paragraph.title}", ) await db.commit() document.status = "completed" if failed_count == 0 else "failed" document.error = "" if failed_count == 0 else f"{failed_count} 个段落生成失败，已回退为模拟结果。" document.file_path = f"mock://document/{document.id}" document.updated_at = datetime.now() await db.commit() update_progress( document_id, status=document.status, percent=100, done=done_count, total=total, message="生成完成" if failed_count == 0 else document.error, ) except Exception as error: document.status = "failed" document.error = str(error) await db.commit() update_progress(document_id, status="failed", message=str(error), done=done_count, total=total) finally: generation_cancel_flags.pop(document_id, None) def build_mock_content_for_block(block: TemplateBlock) -> dict: """为块生成模拟内容。""" if block.output_format == "table": return { "content": [ { "type": "table", "title": block.title, "headers": ["字段", "内容"], "rows": [ ["块标题", block.title], ["块类型", block.block_type], ], } ] } return { "content": [ {"type": "text", "text": f"这是“{block.title or block.block_type}”的示例生成内容。"} ] } async def run_generation_for_blocks(document_id: int, template_id: int): """基于块结构的 AI 生成任务。""" async with async_session() as db: document = await db.get(Document, document_id) template = await db.get(Template, template_id) if document is None or template is None: update_progress(document_id, status="failed", message="生成任务初始化失败") return result = await db.execute( select(TemplateBlock) .where(TemplateBlock.template_id == template_id) .order_by(TemplateBlock.sort_index.asc(), TemplateBlock.id.asc()) ) blocks = result.scalars().all() total = len(blocks) ai_blocks = [b for b in blocks if b.edit_mode == "ai"] total_ai = len(ai_blocks) update_progress(document_id, status="generating", total=total_ai, done=0, percent=0, message="开始块级生成...") done_count = 0 failed_count = 0 # 收集固定块内容作为上下文 fixed_context_parts: list[str] = [] for b in blocks: if b.edit_mode == "manual" and b.title: fixed_context_parts.append(f"[{b.block_type}] {b.title}") try: for block in blocks: if is_cancel_requested(document_id): document.status = "cancelled" document.error = "用户已取消生成" await db.commit() update_progress(document_id, status="cancelled", message="已取消生成", done=done_count) return if block.edit_mode == "manual": # 固定块：使用现有内容 existing = json.loads(block.content_json or "{}") content = {"content": [{"type": "text", "text": existing.get("text", block.title or "")}]} status = "success" duration = 0 error_message = "" model_id = block.model_id else: # AI 块：调用模型生成 start = time.perf_counter() model = None if block.model_id: model = await db.get(AiModel, block.model_id) if model is None or model.status != "enabled": # 尝试使用默认启用的模型 result_enabled = await db.execute( select(AiModel).where(AiModel.status == "enabled").order_by(AiModel.id.asc()).limit(1) ) model = result_enabled.scalars().first() model_id = model.id if model is not None else block.model_id try: if model is None: content = build_mock_content_for_block(block) else: # 构造临时 paragraph 对象用于兼容 call_ai temp_para = Paragraph( template_id=template_id, sort_index=block.sort_index, title=block.title, content="\n".join(fixed_context_parts), prompt_text=block.prompt_text or "", output_format=block.output_format, need_prompt=block.need_prompt, edit_mode=block.edit_mode, ) setattr(temp_para, "enable_reasoning", bool(model.enable_reasoning)) result_data = await call_ai(temp_para, model, None) content = result_data.content status = "success" error_message = "" except Exception as error: content = build_mock_content_for_block(block) status = "failed" error_message = str(error) failed_count += 1 duration = round(time.perf_counter() - start, 4) done_count += 1 log = GenerationLog( document_id=document.id, block_id=block.id, model_id=model_id, status=status, content=json.dumps(content, ensure_ascii=False), duration=duration, error_msg=error_message, ) db.add(log) document.para_count_done = done_count percent = int(done_count / max(total_ai, 1) * 100) update_progress( document_id, status="generating", percent=percent, done=done_count, total=total_ai, current_paragraph=block.title, message=f"正在生成：{block.title or block.block_type}", ) await db.commit() document.status = "completed" if failed_count == 0 else "failed" document.error = "" if failed_count == 0 else f"{failed_count} 个 AI 块生成失败。" document.file_path = f"mock://document/{document.id}" document.updated_at = datetime.now() await db.commit() update_progress( document_id, status=document.status, percent=100, done=done_count, total=total_ai, message="块级生成完成" if failed_count == 0 else document.error, ) except Exception as error: document.status = "failed" document.error = str(error) await db.commit() update_progress(document_id, status="failed", message=str(error), done=done_count, total=total_ai) finally: generation_cancel_flags.pop(document_id, None)