fix: 样式丢失问题 - 生成任务改为直改原始 docx 保留样式
- generate.py: 直接加载原始 docx 用 python-docx 替换文本，不再经过 HTML→docx - html_to_docx: 增强支持 bold/italic/heading 等基本格式 - text替换: 支持段落内 runs 和表格 cell 内的文本替换 - 结果: AI 内容正确插入，字体/加粗/标题等原始样式完全保留
This commit is contained in:
parent
3db8208c55
commit
e937449e6a
|
|
@ -1,7 +1,8 @@
|
|||
import mammoth
|
||||
from io import BytesIO
|
||||
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
|
||||
|
||||
|
||||
|
|
@ -11,6 +12,8 @@ async def docx_to_html(file_content: bytes) -> str:
|
|||
|
||||
|
||||
def html_to_docx_bytes(html_content: str) -> bytes:
|
||||
from html.parser import HTMLParser
|
||||
|
||||
doc = Document()
|
||||
|
||||
style = doc.styles["Normal"]
|
||||
|
|
@ -18,35 +21,75 @@ def html_to_docx_bytes(html_content: str) -> bytes:
|
|||
font.name = "Arial"
|
||||
font.size = Pt(11)
|
||||
|
||||
from html.parser import HTMLParser
|
||||
|
||||
class DocxHTMLParser(HTMLParser):
|
||||
class RichHTMLParser(HTMLParser):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.in_p = False
|
||||
self.current_text = ""
|
||||
self.paragraphs: list[str] = []
|
||||
self.paragraphs: list[dict] = []
|
||||
self.current = {"runs": []}
|
||||
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):
|
||||
if tag in ("p", "h1", "h2", "h3", "h4", "h5", "h6", "div", "li"):
|
||||
self.in_p = True
|
||||
self.current_text = ""
|
||||
tag_lower = tag.lower()
|
||||
if tag_lower in ("p", "div", "li"):
|
||||
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):
|
||||
if tag in ("p", "h1", "h2", "h3", "h4", "h5", "h6", "div", "li"):
|
||||
if self.current_text.strip():
|
||||
self.paragraphs.append(self.current_text.strip())
|
||||
self.current_text = ""
|
||||
self.in_p = False
|
||||
tag_lower = tag.lower()
|
||||
if tag_lower in ("p", "div", "li", "h1", "h2", "h3", "h4", "h5", "h6"):
|
||||
if self.current_run["text"].strip():
|
||||
self.current["runs"].append(dict(self.current_run))
|
||||
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):
|
||||
self.current_text += data
|
||||
if self.in_paragraph:
|
||||
self.current_run["text"] += data
|
||||
|
||||
parser = DocxHTMLParser()
|
||||
parser = RichHTMLParser()
|
||||
parser.feed(html_content)
|
||||
|
||||
for text in parser.paragraphs:
|
||||
doc.add_paragraph(text)
|
||||
for para_data in parser.paragraphs:
|
||||
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()
|
||||
doc.save(output)
|
||||
|
|
|
|||
|
|
@ -1,16 +1,17 @@
|
|||
import os
|
||||
from io import BytesIO
|
||||
from datetime import datetime, timezone
|
||||
from sqlalchemy import select
|
||||
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.core.config import get_settings
|
||||
from app.models.generation_task import GenerationTask
|
||||
from app.models.generation_point import GenerationPoint
|
||||
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.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]:
|
||||
|
|
@ -25,6 +26,79 @@ def _create_db_session() -> async_sessionmaker[AsyncSession]:
|
|||
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:
|
||||
session_factory = _create_db_session()
|
||||
|
||||
|
|
@ -46,6 +120,16 @@ async def _generate_document(task_id: str) -> None:
|
|||
await db.commit()
|
||||
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 ""
|
||||
|
||||
points_result = await db.execute(
|
||||
|
|
@ -57,9 +141,16 @@ async def _generate_document(task_id: str) -> None:
|
|||
|
||||
try:
|
||||
for point in points:
|
||||
selected_text = _get_selected_text(html_content, point.position)
|
||||
if not selected_text:
|
||||
continue
|
||||
|
||||
ref_content = None
|
||||
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": {}}
|
||||
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)
|
||||
|
||||
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:]
|
||||
_replace_text_in_docx(doc, selected_text, ai_result)
|
||||
|
||||
await db.commit()
|
||||
|
||||
result_dir = get_storage_dir(RESULTS_DIR)
|
||||
result_path = os.path.join(result_dir, f"{task_id}.docx")
|
||||
docx_bytes = html_to_docx_bytes(html_content)
|
||||
with open(result_path, "wb") as f:
|
||||
f.write(docx_bytes)
|
||||
doc.save(result_path)
|
||||
|
||||
task.status = "done"
|
||||
task.result_file_path = result_path
|
||||
|
|
|
|||
Loading…
Reference in New Issue