doc-forge/backend/services/generation_runtime.py

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 services.ai_service import call_ai

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)
        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:
                        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)
                            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)