zwt13703
574 lines
21 KiB
Python
574 lines
21 KiB
Python
import asyncio
|
|
import json
|
|
import os
|
|
import tempfile
|
|
import uuid
|
|
from datetime import datetime
|
|
from io import BytesIO
|
|
|
|
from fastapi import APIRouter, Depends, File, HTTPException, Query, UploadFile
|
|
from sqlalchemy import delete, func, select
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from config import settings
|
|
from database import get_db
|
|
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 schemas.schemas import Response, TemplateSave
|
|
from services.document_export import export_document_bytes
|
|
from services.minio_client import download_object_bytes, minio_client, split_bucket_path, upload_bytes
|
|
from services.template_parser import parse_template
|
|
|
|
router = APIRouter()
|
|
|
|
|
|
def _build_object_path(filename: str) -> tuple[str, str]:
|
|
ext = os.path.splitext(filename)[1].lower()
|
|
date_prefix = datetime.now().strftime("%Y%m%d")
|
|
object_name = f"{date_prefix}/{uuid.uuid4().hex}{ext}"
|
|
return ext, object_name
|
|
|
|
|
|
def _serialize_paragraph(paragraph: Paragraph) -> dict:
|
|
return {
|
|
"id": paragraph.id,
|
|
"template_id": paragraph.template_id,
|
|
"sort_index": paragraph.sort_index,
|
|
"anchor_title": paragraph.anchor_title,
|
|
"title": paragraph.title,
|
|
"content": paragraph.content,
|
|
"style_json": paragraph.style_json,
|
|
"is_table": paragraph.is_table,
|
|
"table_json": paragraph.table_json,
|
|
"edit_mode": paragraph.edit_mode,
|
|
"write_mode": paragraph.write_mode,
|
|
"model_id": paragraph.model_id,
|
|
"need_prompt": paragraph.need_prompt,
|
|
"prompt_text": paragraph.prompt_text,
|
|
"need_file": paragraph.need_file,
|
|
"file_note": paragraph.file_note,
|
|
"output_format": paragraph.output_format,
|
|
}
|
|
|
|
|
|
def _serialize_template(template: Template) -> dict:
|
|
return {
|
|
"id": template.id,
|
|
"name": template.name,
|
|
"description": template.description,
|
|
"file_path": template.file_path,
|
|
"paragraph_count": template.paragraph_count,
|
|
"status": template.status,
|
|
"created_at": template.created_at,
|
|
"updated_at": template.updated_at,
|
|
}
|
|
|
|
|
|
def _build_block_from_paragraph(paragraph: Paragraph) -> dict:
|
|
block_type = "heading" if paragraph.write_mode == "replace_heading_only" else ("ai_slot" if paragraph.edit_mode == "ai" else ("table" if paragraph.is_table else "text"))
|
|
content_json = json.dumps({
|
|
"text": paragraph.content or "",
|
|
"table": json.loads(paragraph.table_json or "{}") if paragraph.is_table else None,
|
|
}, ensure_ascii=False)
|
|
return {
|
|
"source_paragraph_id": paragraph.id,
|
|
"parent_block_id": None,
|
|
"sort_index": paragraph.sort_index,
|
|
"block_type": block_type,
|
|
"anchor_ref": paragraph.anchor_title or paragraph.title,
|
|
"title": paragraph.title,
|
|
"content_json": content_json,
|
|
"style_json": paragraph.style_json or "{}",
|
|
"edit_mode": paragraph.edit_mode,
|
|
"placeholder_key": "",
|
|
"variable_key": "",
|
|
"default_value": paragraph.content or "",
|
|
"model_id": paragraph.model_id,
|
|
"need_prompt": paragraph.need_prompt,
|
|
"prompt_text": paragraph.prompt_text,
|
|
"need_file": paragraph.need_file,
|
|
"file_note": paragraph.file_note,
|
|
"output_format": paragraph.output_format,
|
|
}
|
|
|
|
|
|
def _build_block_from_parsed_item(item, source_paragraph_id: int | None) -> dict:
|
|
content_json = json.dumps({
|
|
"text": item.content or "",
|
|
"table": json.loads(item.table_json or "{}") if item.is_table else None,
|
|
}, ensure_ascii=False)
|
|
return {
|
|
"source_paragraph_id": source_paragraph_id,
|
|
"parent_block_id": None,
|
|
"sort_index": item.sort_index,
|
|
"block_type": item.block_type,
|
|
"anchor_ref": item.anchor_title or item.title,
|
|
"title": item.title,
|
|
"content_json": content_json,
|
|
"style_json": item.style_json or "{}",
|
|
"edit_mode": item.edit_mode,
|
|
"placeholder_key": item.placeholder_key,
|
|
"variable_key": item.variable_key,
|
|
"default_value": item.default_value,
|
|
"model_id": None,
|
|
"need_prompt": True,
|
|
"prompt_text": "",
|
|
"need_file": False,
|
|
"file_note": "",
|
|
"output_format": item.output_format,
|
|
}
|
|
|
|
|
|
def _serialize_block(block: TemplateBlock) -> dict:
|
|
try:
|
|
content_json = json.loads(block.content_json or "{}")
|
|
except Exception:
|
|
content_json = {}
|
|
return {
|
|
"id": block.id,
|
|
"template_id": block.template_id,
|
|
"source_paragraph_id": block.source_paragraph_id,
|
|
"parent_block_id": block.parent_block_id,
|
|
"sort_index": block.sort_index,
|
|
"block_type": block.block_type,
|
|
"anchor_ref": block.anchor_ref,
|
|
"title": block.title,
|
|
"content_json": content_json,
|
|
"style_json": block.style_json,
|
|
"edit_mode": block.edit_mode,
|
|
"placeholder_key": block.placeholder_key,
|
|
"variable_key": block.variable_key,
|
|
"default_value": block.default_value,
|
|
"model_id": block.model_id,
|
|
"need_prompt": block.need_prompt,
|
|
"prompt_text": block.prompt_text,
|
|
"need_file": block.need_file,
|
|
"file_note": block.file_note,
|
|
"output_format": block.output_format,
|
|
}
|
|
|
|
|
|
async def _load_blocks(db: AsyncSession, template_id: int) -> list[TemplateBlock]:
|
|
result = await db.execute(
|
|
select(TemplateBlock)
|
|
.where(TemplateBlock.template_id == template_id)
|
|
.order_by(TemplateBlock.sort_index.asc(), TemplateBlock.id.asc())
|
|
)
|
|
return result.scalars().all()
|
|
|
|
|
|
async def _sync_blocks_from_paragraphs(db: AsyncSession, template_id: int, paragraphs: list[Paragraph]):
|
|
existing_blocks = await _load_blocks(db, template_id)
|
|
for block in existing_blocks:
|
|
await db.delete(block)
|
|
await db.flush()
|
|
|
|
block_rows: list[TemplateBlock] = []
|
|
for paragraph in paragraphs:
|
|
block = TemplateBlock(template_id=template_id, **_build_block_from_paragraph(paragraph))
|
|
db.add(block)
|
|
block_rows.append(block)
|
|
await db.flush()
|
|
return block_rows
|
|
|
|
|
|
async def _save_blocks(
|
|
db: AsyncSession,
|
|
template_id: int,
|
|
blocks_payload,
|
|
):
|
|
existing_blocks = await _load_blocks(db, template_id)
|
|
block_map = {item.id: item for item in existing_blocks}
|
|
incoming_ids = {config.id for config in blocks_payload if config.id}
|
|
|
|
for block in existing_blocks:
|
|
if block.id not in incoming_ids:
|
|
await db.delete(block)
|
|
|
|
for index, config in enumerate(blocks_payload, start=1):
|
|
block = block_map.get(config.id) if config.id else None
|
|
if block is None:
|
|
block = TemplateBlock(template_id=template_id)
|
|
db.add(block)
|
|
block.source_paragraph_id = config.source_paragraph_id
|
|
block.parent_block_id = config.parent_block_id
|
|
block.sort_index = index
|
|
block.block_type = config.block_type
|
|
block.anchor_ref = config.anchor_ref or config.title
|
|
block.title = config.title
|
|
block.content_json = json.dumps(config.content_json or {}, ensure_ascii=False)
|
|
block.style_json = config.style_json or "{}"
|
|
block.edit_mode = config.edit_mode
|
|
block.placeholder_key = config.placeholder_key
|
|
block.variable_key = config.variable_key
|
|
block.default_value = config.default_value
|
|
block.model_id = config.model_id
|
|
block.need_prompt = config.need_prompt
|
|
block.prompt_text = config.prompt_text
|
|
block.need_file = config.need_file
|
|
block.file_note = config.file_note
|
|
block.output_format = config.output_format
|
|
|
|
await db.flush()
|
|
|
|
|
|
def _block_text_content(block: TemplateBlock) -> str:
|
|
try:
|
|
payload = json.loads(block.content_json or "{}")
|
|
except Exception:
|
|
payload = {}
|
|
return payload.get("text") or block.default_value or ""
|
|
|
|
|
|
def _block_table_content(block: TemplateBlock) -> dict:
|
|
try:
|
|
payload = json.loads(block.content_json or "{}")
|
|
except Exception:
|
|
payload = {}
|
|
return payload.get("table") or {}
|
|
|
|
|
|
async def _sync_paragraphs_from_blocks(db: AsyncSession, template_id: int) -> list[Paragraph]:
|
|
paragraph_result = await db.execute(
|
|
select(Paragraph)
|
|
.where(Paragraph.template_id == template_id)
|
|
.order_by(Paragraph.sort_index.asc(), Paragraph.id.asc())
|
|
)
|
|
existing_paragraphs = paragraph_result.scalars().all()
|
|
paragraph_map = {item.id: item for item in existing_paragraphs}
|
|
blocks = await _load_blocks(db, template_id)
|
|
write_modes = _resolve_block_write_modes(blocks)
|
|
kept_paragraph_ids: set[int] = set()
|
|
synced_rows: list[Paragraph] = []
|
|
|
|
for index, (block, write_mode) in enumerate(zip(blocks, write_modes), start=1):
|
|
paragraph = paragraph_map.get(block.source_paragraph_id) if block.source_paragraph_id else None
|
|
if paragraph is None:
|
|
paragraph = Paragraph(template_id=template_id)
|
|
db.add(paragraph)
|
|
await db.flush()
|
|
paragraph.sort_index = index
|
|
paragraph.anchor_title = block.anchor_ref or block.title
|
|
paragraph.title = block.title
|
|
paragraph.content = _block_text_content(block)
|
|
paragraph.style_json = block.style_json or "{}"
|
|
paragraph.is_table = block.block_type == "table"
|
|
paragraph.table_json = json.dumps(_block_table_content(block), ensure_ascii=False) if paragraph.is_table else "{}"
|
|
paragraph.edit_mode = "ai" if block.block_type == "ai_slot" or block.edit_mode == "ai" else "manual"
|
|
paragraph.write_mode = write_mode
|
|
paragraph.model_id = block.model_id
|
|
paragraph.need_prompt = block.need_prompt
|
|
paragraph.prompt_text = block.prompt_text
|
|
paragraph.need_file = block.need_file
|
|
paragraph.file_note = block.file_note
|
|
paragraph.output_format = block.output_format
|
|
block.source_paragraph_id = paragraph.id
|
|
kept_paragraph_ids.add(paragraph.id)
|
|
synced_rows.append(paragraph)
|
|
|
|
for paragraph in existing_paragraphs:
|
|
if paragraph.id in kept_paragraph_ids:
|
|
continue
|
|
await db.execute(delete(GenerationLog).where(GenerationLog.paragraph_id == paragraph.id))
|
|
await db.delete(paragraph)
|
|
|
|
await db.flush()
|
|
return synced_rows
|
|
|
|
|
|
def _build_export_content_from_block(block: TemplateBlock) -> dict:
|
|
if block.block_type == "table":
|
|
table_data = _block_table_content(block)
|
|
matrix = table_data.get("data") or []
|
|
headers = matrix[0] if matrix else []
|
|
rows = matrix[1:] if len(matrix) > 1 else []
|
|
return {"content": [{"type": "table", "headers": headers, "rows": rows}]}
|
|
return {"content": [{"type": "text", "text": _block_text_content(block)}]}
|
|
|
|
|
|
def _resolve_block_write_modes(blocks: list[TemplateBlock]) -> list[str]:
|
|
modes: list[str] = []
|
|
anchor_counter: dict[str, int] = {}
|
|
for block in blocks:
|
|
if block.block_type == "heading":
|
|
modes.append("replace_heading_only")
|
|
continue
|
|
anchor = (block.anchor_ref or block.title or "").strip()
|
|
seen = anchor_counter.get(anchor, 0)
|
|
modes.append("replace_section" if seen == 0 else "append_after_heading")
|
|
anchor_counter[anchor] = seen + 1
|
|
return modes
|
|
|
|
|
|
async def _write_template_snapshot_to_docx(db: AsyncSession, template: Template):
|
|
blocks = await _load_blocks(db, template.id)
|
|
if not blocks:
|
|
return
|
|
write_modes = _resolve_block_write_modes(blocks)
|
|
logs = []
|
|
for block, write_mode in zip(blocks, write_modes):
|
|
logs.append(
|
|
{
|
|
"anchor_title": block.anchor_ref or block.title,
|
|
"title": block.title,
|
|
"write_mode": write_mode,
|
|
"content": _build_export_content_from_block(block),
|
|
}
|
|
)
|
|
|
|
template_bucket, template_object = split_bucket_path(template.file_path)
|
|
template_bytes = await asyncio.to_thread(download_object_bytes, template_bucket, template_object)
|
|
exported_bytes = await asyncio.to_thread(export_document_bytes, template_bytes, logs)
|
|
await asyncio.to_thread(
|
|
upload_bytes,
|
|
template_bucket,
|
|
template_object,
|
|
exported_bytes,
|
|
"application/vnd.openxmlformats-officedocument.wordprocessingml.document",
|
|
)
|
|
|
|
|
|
@router.get("")
|
|
async def list_templates(
|
|
page: int = Query(1, ge=1),
|
|
page_size: int = Query(20, ge=1, le=100),
|
|
keyword: str = Query("", alias="q"),
|
|
db: AsyncSession = Depends(get_db),
|
|
):
|
|
filters = []
|
|
if keyword:
|
|
filters.append(Template.name.like(f"%{keyword}%"))
|
|
|
|
total_stmt = select(func.count(Template.id))
|
|
list_stmt = select(Template).order_by(Template.id.desc())
|
|
if filters:
|
|
total_stmt = total_stmt.where(*filters)
|
|
list_stmt = list_stmt.where(*filters)
|
|
|
|
total = (await db.execute(total_stmt)).scalar_one()
|
|
result = await db.execute(list_stmt.offset((page - 1) * page_size).limit(page_size))
|
|
items = [_serialize_template(item) for item in result.scalars().all()]
|
|
return Response(
|
|
data={"items": items, "total": total, "page": page, "page_size": page_size}
|
|
)
|
|
|
|
|
|
@router.get("/{template_id}")
|
|
async def get_template(template_id: int, db: AsyncSession = Depends(get_db)):
|
|
template = await db.get(Template, template_id)
|
|
if template is None:
|
|
raise HTTPException(status_code=404, detail="模板不存在")
|
|
|
|
result = await db.execute(
|
|
select(Paragraph)
|
|
.where(Paragraph.template_id == template_id)
|
|
.order_by(Paragraph.sort_index.asc(), Paragraph.id.asc())
|
|
)
|
|
paragraph_rows = result.scalars().all()
|
|
paragraphs = [_serialize_paragraph(item) for item in paragraph_rows]
|
|
block_rows = await _load_blocks(db, template_id)
|
|
if not block_rows and paragraph_rows:
|
|
block_rows = await _sync_blocks_from_paragraphs(db, template_id, paragraph_rows)
|
|
await db.commit()
|
|
blocks = [_serialize_block(item) for item in block_rows]
|
|
payload = _serialize_template(template)
|
|
payload["paragraphs"] = paragraphs
|
|
payload["blocks"] = blocks
|
|
return Response(data=payload)
|
|
|
|
|
|
@router.post("/upload")
|
|
async def upload_template(file: UploadFile = File(...), db: AsyncSession = Depends(get_db)):
|
|
if not file.filename:
|
|
raise HTTPException(status_code=400, detail="文件名不能为空")
|
|
|
|
ext, object_name = _build_object_path(file.filename)
|
|
if ext != ".docx":
|
|
raise HTTPException(status_code=400, detail="模板仅支持 .docx 格式")
|
|
|
|
content = await file.read()
|
|
if not content:
|
|
raise HTTPException(status_code=400, detail="上传文件不能为空")
|
|
if len(content) > settings.MAX_UPLOAD_SIZE:
|
|
raise HTTPException(status_code=400, detail="文件大小超过限制")
|
|
|
|
with tempfile.NamedTemporaryFile(delete=False, suffix=ext) as temp_file:
|
|
temp_file.write(content)
|
|
temp_path = temp_file.name
|
|
|
|
try:
|
|
parsed_items = await asyncio.to_thread(parse_template, temp_path)
|
|
finally:
|
|
if os.path.exists(temp_path):
|
|
os.remove(temp_path)
|
|
|
|
await asyncio.to_thread(
|
|
minio_client.put_object,
|
|
settings.MINIO_BUCKET_TEMPLATES,
|
|
object_name,
|
|
BytesIO(content),
|
|
len(content),
|
|
file.content_type or "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
|
|
)
|
|
|
|
template = Template(
|
|
name=os.path.splitext(file.filename)[0],
|
|
description="",
|
|
file_path=f"{settings.MINIO_BUCKET_TEMPLATES}/{object_name}",
|
|
paragraph_count=len(parsed_items),
|
|
status="draft",
|
|
)
|
|
db.add(template)
|
|
await db.flush()
|
|
|
|
paragraph_rows: list[Paragraph] = []
|
|
for item in parsed_items:
|
|
paragraph = Paragraph(
|
|
template_id=template.id,
|
|
sort_index=item.sort_index,
|
|
anchor_title=item.anchor_title,
|
|
title=item.title,
|
|
content=item.content,
|
|
style_json=item.style_json,
|
|
is_table=item.is_table,
|
|
table_json=item.table_json,
|
|
edit_mode=item.edit_mode,
|
|
write_mode=item.write_mode,
|
|
need_prompt=item.edit_mode == "ai",
|
|
output_format=item.output_format,
|
|
)
|
|
db.add(paragraph)
|
|
paragraph_rows.append(paragraph)
|
|
|
|
await db.flush()
|
|
existing_blocks = await _load_blocks(db, template.id)
|
|
for block in existing_blocks:
|
|
await db.delete(block)
|
|
await db.flush()
|
|
block_rows: list[TemplateBlock] = []
|
|
for item, paragraph in zip(parsed_items, paragraph_rows):
|
|
block = TemplateBlock(template_id=template.id, **_build_block_from_parsed_item(item, paragraph.id))
|
|
db.add(block)
|
|
block_rows.append(block)
|
|
await db.flush()
|
|
await db.commit()
|
|
await db.refresh(template)
|
|
for paragraph in paragraph_rows:
|
|
await db.refresh(paragraph)
|
|
for block in block_rows:
|
|
await db.refresh(block)
|
|
|
|
payload = _serialize_template(template)
|
|
payload["paragraphs"] = [_serialize_paragraph(item) for item in paragraph_rows]
|
|
payload["blocks"] = [_serialize_block(item) for item in block_rows]
|
|
return Response(data=payload)
|
|
|
|
|
|
@router.put("/{template_id}/paragraphs")
|
|
async def save_template_paragraphs(
|
|
template_id: int,
|
|
body: TemplateSave,
|
|
db: AsyncSession = Depends(get_db),
|
|
):
|
|
template = await db.get(Template, template_id)
|
|
if template is None:
|
|
raise HTTPException(status_code=404, detail="模板不存在")
|
|
|
|
result = await db.execute(
|
|
select(Paragraph)
|
|
.where(Paragraph.template_id == template_id)
|
|
.order_by(Paragraph.sort_index.asc(), Paragraph.id.asc())
|
|
)
|
|
existing_paragraphs = result.scalars().all()
|
|
paragraph_map = {item.id: item for item in existing_paragraphs}
|
|
incoming_ids = {config.id for config in body.paragraphs if config.id}
|
|
|
|
print(f"[SAVE] template_id={template_id}, incoming_ids={incoming_ids}, existing_ids={[p.id for p in existing_paragraphs]}")
|
|
|
|
for paragraph in existing_paragraphs:
|
|
if paragraph.id not in incoming_ids:
|
|
print(f"[SAVE] Deleting paragraph id={paragraph.id} title={paragraph.title}")
|
|
await db.execute(
|
|
delete(GenerationLog).where(GenerationLog.paragraph_id == paragraph.id)
|
|
)
|
|
await db.delete(paragraph)
|
|
|
|
for index, config in enumerate(body.paragraphs, start=1):
|
|
paragraph = paragraph_map.get(config.id) if config.id else None
|
|
if paragraph is None:
|
|
paragraph = Paragraph(template_id=template_id)
|
|
db.add(paragraph)
|
|
paragraph.sort_index = index
|
|
paragraph.anchor_title = config.anchor_title or config.title or paragraph.anchor_title
|
|
paragraph.title = config.title
|
|
paragraph.content = config.content
|
|
paragraph.edit_mode = config.edit_mode
|
|
paragraph.write_mode = config.write_mode
|
|
paragraph.model_id = config.model_id
|
|
paragraph.need_prompt = config.need_prompt
|
|
paragraph.prompt_text = config.prompt_text
|
|
paragraph.need_file = config.need_file
|
|
paragraph.file_note = config.file_note
|
|
paragraph.output_format = config.output_format
|
|
|
|
await db.flush()
|
|
refreshed_result = await db.execute(
|
|
select(Paragraph)
|
|
.where(Paragraph.template_id == template_id)
|
|
.order_by(Paragraph.sort_index.asc(), Paragraph.id.asc())
|
|
)
|
|
refreshed_paragraphs = refreshed_result.scalars().all()
|
|
if body.save_mode == "manual" and body.blocks:
|
|
await _save_blocks(db, template_id, body.blocks)
|
|
refreshed_paragraphs = await _sync_paragraphs_from_blocks(db, template_id)
|
|
else:
|
|
await _sync_blocks_from_paragraphs(db, template_id, refreshed_paragraphs)
|
|
refreshed_paragraphs = await _sync_paragraphs_from_blocks(db, template_id)
|
|
template.paragraph_count = len(refreshed_paragraphs)
|
|
await _write_template_snapshot_to_docx(db, template)
|
|
await db.commit()
|
|
blocks = [_serialize_block(item) for item in await _load_blocks(db, template_id)]
|
|
return Response(data={"template_id": template_id, "saved": len(refreshed_paragraphs), "blocks": blocks})
|
|
|
|
|
|
@router.delete("/{template_id}")
|
|
async def delete_template(template_id: int, db: AsyncSession = Depends(get_db)):
|
|
template = await db.get(Template, template_id)
|
|
if template is None:
|
|
raise HTTPException(status_code=404, detail="模板不存在")
|
|
|
|
result = await db.execute(select(Paragraph).where(Paragraph.template_id == template_id))
|
|
paragraphs_to_delete = result.scalars().all()
|
|
paragraph_ids = [p.id for p in paragraphs_to_delete]
|
|
|
|
doc_result = await db.execute(select(Document).where(Document.template_id == template_id))
|
|
documents_to_delete = doc_result.scalars().all()
|
|
|
|
if paragraph_ids:
|
|
await db.execute(
|
|
delete(GenerationLog).where(GenerationLog.paragraph_id.in_(paragraph_ids))
|
|
)
|
|
for document in documents_to_delete:
|
|
await db.execute(
|
|
delete(GenerationLog).where(GenerationLog.document_id == document.id)
|
|
)
|
|
await db.delete(document)
|
|
|
|
for paragraph in paragraphs_to_delete:
|
|
await db.delete(paragraph)
|
|
|
|
file_path = template.file_path or ""
|
|
if "/" in file_path:
|
|
bucket, object_name = file_path.split("/", 1)
|
|
try:
|
|
await asyncio.to_thread(minio_client.remove_object, bucket, object_name)
|
|
except Exception:
|
|
pass
|
|
|
|
await db.delete(template)
|
|
await db.commit()
|
|
return Response(data={"id": template_id})
|