diff --git a/backend/app/services/document_processor.py b/backend/app/services/document_processor.py index dc34ddf..65d88ae 100644 --- a/backend/app/services/document_processor.py +++ b/backend/app/services/document_processor.py @@ -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) diff --git a/backend/app/tasks/generate.py b/backend/app/tasks/generate.py index b3aeb10..685c77c 100644 --- a/backend/app/tasks/generate.py +++ b/backend/app/tasks/generate.py @@ -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