按原型图重构核心页面并完善生成链路
This commit is contained in:
zwt13703
@@ -0,0 +1,174 @@
|
||||
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:
|
||||
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)
|
||||
Reference in New Issue
Block a user