fix: 样式丢失问题 - 生成任务改为直改原始 docx 保留样式

- generate.py: 直接加载原始 docx 用 python-docx 替换文本,不再经过 HTML→docx
- html_to_docx: 增强支持 bold/italic/heading 等基本格式
- text替换: 支持段落内 runs 和表格 cell 内的文本替换
- 结果: AI 内容正确插入,字体/加粗/标题等原始样式完全保留
This commit is contained in:
zwt13703 2026-07-06 17:53:15 +08:00
parent 3db8208c55
commit e937449e6a
2 changed files with 158 additions and 31 deletions

View File

@ -1,7 +1,8 @@
import mammoth import mammoth
from io import BytesIO from io import BytesIO
from docx import Document from docx import Document
from docx.shared import Pt from docx.shared import Pt, Inches
from docx.enum.text import WD_ALIGN_PARAGRAPH
from app.services.file_storage import get_file_content from app.services.file_storage import get_file_content
@ -11,6 +12,8 @@ async def docx_to_html(file_content: bytes) -> str:
def html_to_docx_bytes(html_content: str) -> bytes: def html_to_docx_bytes(html_content: str) -> bytes:
from html.parser import HTMLParser
doc = Document() doc = Document()
style = doc.styles["Normal"] style = doc.styles["Normal"]
@ -18,35 +21,75 @@ def html_to_docx_bytes(html_content: str) -> bytes:
font.name = "Arial" font.name = "Arial"
font.size = Pt(11) font.size = Pt(11)
from html.parser import HTMLParser class RichHTMLParser(HTMLParser):
class DocxHTMLParser(HTMLParser):
def __init__(self): def __init__(self):
super().__init__() super().__init__()
self.in_p = False self.paragraphs: list[dict] = []
self.current_text = "" self.current = {"runs": []}
self.paragraphs: list[str] = [] self.in_paragraph = False
self.current_run = {"text": "", "bold": False, "italic": False, "underline": False}
self.tag_stack: list[str] = []
self.heading_level = 0
def handle_starttag(self, tag, attrs): def handle_starttag(self, tag, attrs):
if tag in ("p", "h1", "h2", "h3", "h4", "h5", "h6", "div", "li"): tag_lower = tag.lower()
self.in_p = True if tag_lower in ("p", "div", "li"):
self.current_text = "" self.in_paragraph = True
self.current_run = {"text": "", "bold": False, "italic": False, "underline": False}
elif tag_lower in ("h1", "h2", "h3", "h4", "h5", "h6"):
self.in_paragraph = True
self.heading_level = int(tag_lower[1])
self.current_run = {"text": "", "bold": False, "italic": False, "underline": False}
elif tag_lower in ("strong", "b"):
self.current_run["bold"] = True
elif tag_lower in ("em", "i"):
self.current_run["italic"] = True
elif tag_lower == "u":
self.current_run["underline"] = True
elif tag_lower in ("br",):
if self.in_paragraph:
self.current["runs"].append(dict(self.current_run))
self.current_run = {"text": "", "bold": False, "italic": False, "underline": False}
self.tag_stack.append(tag_lower)
def handle_endtag(self, tag): def handle_endtag(self, tag):
if tag in ("p", "h1", "h2", "h3", "h4", "h5", "h6", "div", "li"): tag_lower = tag.lower()
if self.current_text.strip(): if tag_lower in ("p", "div", "li", "h1", "h2", "h3", "h4", "h5", "h6"):
self.paragraphs.append(self.current_text.strip()) if self.current_run["text"].strip():
self.current_text = "" self.current["runs"].append(dict(self.current_run))
self.in_p = False if self.current["runs"]:
p = dict(self.current)
p["heading"] = self.heading_level
self.paragraphs.append(p)
self.current = {"runs": []}
self.current_run = {"text": "", "bold": False, "italic": False, "underline": False}
self.in_paragraph = False
self.heading_level = 0
if self.tag_stack:
self.tag_stack.pop()
def handle_data(self, data): def handle_data(self, data):
self.current_text += data if self.in_paragraph:
self.current_run["text"] += data
parser = DocxHTMLParser() parser = RichHTMLParser()
parser.feed(html_content) parser.feed(html_content)
for text in parser.paragraphs: for para_data in parser.paragraphs:
doc.add_paragraph(text) heading = para_data.get("heading", 0)
if heading > 0:
p = doc.add_heading(level=min(heading, 9))
else:
p = doc.add_paragraph()
for run_data in para_data.get("runs", []):
run = p.add_run(run_data["text"])
if run_data.get("bold"):
run.bold = True
if run_data.get("italic"):
run.italic = True
if run_data.get("underline"):
run.underline = True
output = BytesIO() output = BytesIO()
doc.save(output) doc.save(output)

View File

@ -1,16 +1,17 @@
import os import os
from io import BytesIO
from datetime import datetime, timezone from datetime import datetime, timezone
from sqlalchemy import select from sqlalchemy import select
from sqlalchemy.ext.asyncio import create_async_engine, async_sessionmaker, AsyncSession from sqlalchemy.ext.asyncio import create_async_engine, async_sessionmaker, AsyncSession
from docx import Document
from app.tasks.celery_app import celery_app from app.tasks.celery_app import celery_app
from app.core.config import get_settings from app.core.config import get_settings
from app.models.generation_task import GenerationTask from app.models.generation_task import GenerationTask
from app.models.generation_point import GenerationPoint from app.models.generation_point import GenerationPoint
from app.models.template import Template from app.models.template import Template
from app.services.document_processor import html_to_docx_bytes
from app.services.ai_adapter import call_ai_model from app.services.ai_adapter import call_ai_model
from app.services.ref_parser import parse_reference_file from app.services.ref_parser import parse_reference_file
from app.services.file_storage import get_storage_dir, RESULTS_DIR from app.services.file_storage import get_storage_dir, get_file_content, RESULTS_DIR
def _create_db_session() -> async_sessionmaker[AsyncSession]: def _create_db_session() -> async_sessionmaker[AsyncSession]:
@ -25,6 +26,79 @@ def _create_db_session() -> async_sessionmaker[AsyncSession]:
return async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False) return async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)
def _get_selected_text(html_content: str, position: dict) -> str:
start = position.get("start", 0)
end = position.get("end", 0)
if 0 <= start < end <= len(html_content):
from html.parser import HTMLParser
class TextStripper(HTMLParser):
def __init__(self):
super().__init__()
self.text = ""
def handle_data(self, data):
self.text += data
stripper = TextStripper()
stripper.feed(html_content[start:end])
return stripper.text.strip()
return ""
def _replace_text_in_docx(doc: Document, old_text: str, new_text: str) -> bool:
if not old_text:
return False
for paragraph in doc.paragraphs:
if old_text in paragraph.text:
inline = paragraph.runs
for run in inline:
if old_text in run.text:
run.text = run.text.replace(old_text, new_text)
return True
full_text = "".join(r.text for r in inline)
if old_text in full_text:
remaining = old_text
for run in inline:
if not remaining:
break
if remaining.startswith(run.text):
remaining = remaining[len(run.text):]
elif run.text in remaining:
idx = remaining.find(run.text)
if idx >= 0:
remaining = remaining[:idx] + remaining[idx + len(run.text):]
if remaining.startswith(run.text):
remaining = remaining[len(run.text):]
if not remaining:
chunk_parts = new_text
for run in inline:
if chunk_parts:
chunk_parts = chunk_parts[len(run.text):]
first_run = inline[0]
first_run.text = new_text
for run in inline[1:]:
run.text = ""
return True
for table in doc.tables:
for row in table.rows:
for cell in row.cells:
for paragraph in cell.paragraphs:
if old_text in paragraph.text:
for run in paragraph.runs:
if old_text in run.text:
run.text = run.text.replace(old_text, new_text)
return True
return False
async def _generate_document(task_id: str) -> None: async def _generate_document(task_id: str) -> None:
session_factory = _create_db_session() session_factory = _create_db_session()
@ -46,6 +120,16 @@ async def _generate_document(task_id: str) -> None:
await db.commit() await db.commit()
return return
if not os.path.exists(template.file_path):
task.status = "failed"
task.error_msg = f"原始文件不存在: {template.file_path}"
task.finished_at = datetime.now(timezone.utc)
await db.commit()
return
docx_content = await get_file_content(template.file_path)
doc = Document(BytesIO(docx_content))
html_content = template.html_content or "" html_content = template.html_content or ""
points_result = await db.execute( points_result = await db.execute(
@ -57,9 +141,16 @@ async def _generate_document(task_id: str) -> None:
try: try:
for point in points: for point in points:
selected_text = _get_selected_text(html_content, point.position)
if not selected_text:
continue
ref_content = None ref_content = None
if point.ref_file_path: if point.ref_file_path:
ref_content = await parse_reference_file(point.ref_file_path) try:
ref_content = await parse_reference_file(point.ref_file_path)
except Exception:
pass
model_config = {"provider": "custom", "endpoint": "", "api_key": "", "extra_params": {}} model_config = {"provider": "custom", "endpoint": "", "api_key": "", "extra_params": {}}
if point.model_id: if point.model_id:
@ -75,20 +166,13 @@ async def _generate_document(task_id: str) -> None:
} }
ai_result = await call_ai_model(model_config, point.prompt, ref_content) ai_result = await call_ai_model(model_config, point.prompt, ref_content)
_replace_text_in_docx(doc, selected_text, ai_result)
position = point.position
start = position.get("start", 0)
end = position.get("end", 0)
if 0 <= start < end <= len(html_content):
html_content = html_content[:start] + ai_result + html_content[end:]
await db.commit() await db.commit()
result_dir = get_storage_dir(RESULTS_DIR) result_dir = get_storage_dir(RESULTS_DIR)
result_path = os.path.join(result_dir, f"{task_id}.docx") result_path = os.path.join(result_dir, f"{task_id}.docx")
docx_bytes = html_to_docx_bytes(html_content) doc.save(result_path)
with open(result_path, "wb") as f:
f.write(docx_bytes)
task.status = "done" task.status = "done"
task.result_file_path = result_path task.result_file_path = result_path