from docx import Document from docx.shared import Pt, Inches from docx.oxml.ns import qn from io import BytesIO from docx.enum.text import WD_ALIGN_PARAGRAPH import random import threading import re import time import json from tenacity import retry, stop_after_attempt, wait_fixed from io import BytesIO from PIL import Image from auto_media_publisher.config.conf_base import GOOGLE_GEMINI_PROXY from auto_media_publisher.config.conf_base import GOOGLE_CHAT_MIN_INTERVAL,DEFAULT_PREFIX_MAP,DEFAULT_SURFIX_MAP from auto_media_publisher.config.conf_base import GOOGLE_REWRITE_STYLE_MAP,GEMINI_PROMPT_INIT_DEFAULT from auto_media_publisher.config.conf_base import GEMINI_PROMPT_TITLE,GEMINI_PROMPT_DESCRIPTION,GEMINI_PROMPT_TAG from auto_media_publisher.utils.logger import get_logger from auto_media_publisher.utils.utils_text2image_pollinations import Text2ImageGenerator from auto_media_publisher.utils.utils_google import init_gemini_model class DocxRewriterGemini: def __init__(self,rewrite_style="default" ,docx_author="我"): self.model = init_gemini_model(GOOGLE_GEMINI_PROXY) if rewrite_style =="random": self.rewrite_style = random.choice(list(GOOGLE_REWRITE_STYLE_MAP.keys())) else: self.rewrite_style = rewrite_style self.prefix_text = f"我是{docx_author}，{random.choice(DEFAULT_PREFIX_MAP)}" self.surfix_text = random.choice(DEFAULT_SURFIX_MAP) self.last_call_time = 0 self.min_interval = GOOGLE_CHAT_MIN_INTERVAL # 限制最短请求间隔（秒） self.lock = threading.Lock() # 多线程并发时也安全 self.logger = get_logger(self.__class__.__name__,"admin") self.text2imager = Text2ImageGenerator(width=900,height=383) self.summary_info = {} #@retry(stop=stop_after_attempt(3), wait=wait_fixed(2)) # 自动重试机制 def safe_send_message(self, content: str): with self.lock: now = time.time() wait_time = self.min_interval - (now - self.last_call_time) if wait_time > 0: time.sleep(wait_time) self.last_call_time = time.time() return self.chat.send_message(content) def rewrite(self, input_path: str, output_path: str): blocks = self.extract_docx_blocks_by_order(input_path) self.logger.info(f"start to rewriting docx using gemini,. with style {self.rewrite_style}. {input_path}") if len(blocks) >0: self.chat = self.model.start_chat(history=[]) self.safe_send_message(GOOGLE_REWRITE_STYLE_MAP.get(self.rewrite_style, GEMINI_PROMPT_INIT_DEFAULT)) rewritten = self._rewrite_blocks(blocks) self._save_to_new_docx(rewritten, output_path) self.summary_info = self._summarize_title_intro_keywords() print(self.summary_info) self.text2imager.generate_image(self.summary_info["description"],output_path.replace(".docx",".png")) self.summary_info["path_cover"] = output_path.replace(".docx",".png") def extract_docx_blocks_by_order(self,docx_path): doc = Document(docx_path) blocks = [] current_texts = [] def _get_blip_rids(run_element): blips = [] for blip in run_element.findall('.//{http://schemas.openxmlformats.org/drawingml/2006/main}blip'): rId = blip.get('{http://schemas.openxmlformats.org/officeDocument/2006/relationships}embed') if rId: blips.append(rId) return blips # 遍历文档的 body 中的元素 for child in doc.element.body: # 处理段落 if child.tag == qn('w:p'): para = next((p for p in doc.paragraphs if p._element == child), None) if para is not None: # 判断此段是否包含图片 has_image = False images = [] for run in para.runs: rIds = _get_blip_rids(run.element) for rId in rIds: try: image_part = run.part.related_parts[rId] images.append(image_part.blob) has_image = True except KeyError: continue if has_image: if current_texts: blocks.append({'type': 'text', 'content': '\n'.join(current_texts)}) current_texts = [] for img in images: blocks.append({'type': 'image', 'content': img}) else: text = para.text.strip() if text: current_texts.append(text) # 结尾残留文本也写入 block if current_texts: blocks.append({'type': 'text', 'content': '\n'.join(current_texts)}) return blocks def _extract_image_bytes(self, para): # 获取图片字节数据 for run in para.runs: drawing = run._element.xpath('.//a:blip') if drawing: rId = drawing[0].get("r:embed") part = run.part.related_parts[rId] return part.blob return None def _rewrite_blocks(self, blocks): rewritten_blocks = [] print("start to send oringal message to gemini ......") for i, block in enumerate(blocks): #print(block) if block["type"]=="text": prompt = f"{chr(12288).join(block['content'])}" response = self.safe_send_message(prompt) rewritten_text = response.text.strip() rewritten_blocks.append({"type": "text", "content": rewritten_text}) else: rewritten_blocks.append(block) return rewritten_blocks def _save_to_new_docx(self, blocks, output_path): doc = Document() # 设置默认样式 style = doc.styles['Normal'] font = style.font font.name = '楷体' # 西文字体设置为楷体 font.size = Pt(10.5) # 小四字体 style._element.rPr.rFonts.set(qn('w:eastAsia'), '楷体') # 中文字体设定 # 前言 para = doc.add_paragraph(self.prefix_text) para.paragraph_format.line_spacing = 1.5 para.paragraph_format.space_before = Pt(0) para.paragraph_format.space_after = Pt(0) for block in blocks: if block["type"] == "text": for line in block["content"].split("\n"): para = doc.add_paragraph(line.strip()) para.paragraph_format.line_spacing = 1.5 # 行距1.5倍 para.paragraph_format.space_before = Pt(0) # 段前0磅 para.paragraph_format.space_after = Pt(0) # 段后0磅 elif block["type"] == "image": image_data = block["content"] image_stream = BytesIO(image_data) paragraph = doc.add_paragraph() run = paragraph.add_run() run.add_picture(image_stream, width=Inches(5.5)) paragraph.alignment = WD_ALIGN_PARAGRAPH.CENTER # 结语 para2 = doc.add_paragraph(self.surfix_text) para2.paragraph_format.line_spacing = 1.5 para2.paragraph_format.space_before = Pt(12) para2.paragraph_format.space_after = Pt(0) doc.save(output_path) def _summarize_title_intro_keywords(self) -> dict: title = self.safe_send_message(GEMINI_PROMPT_TITLE).text.strip() description = self.safe_send_message(GEMINI_PROMPT_DESCRIPTION).text.strip() keywords = self.safe_send_message(GEMINI_PROMPT_TAG).text.strip() return { "title": title.splitlines()[0].strip()[:64], "description": description[:120], "tags": re.split(r'[、,\s]+', keywords.strip("。")) }