356 lines
15 KiB
Python
356 lines
15 KiB
Python
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)
|