from docx import Document from docx.shared import Inches from docx.enum.text import WD_ALIGN_PARAGRAPH from docx.oxml.ns import qn 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 import google.generativeai as genai from auto_media_publisher.config.conf_base import GOOGLE_API_KEY,GOOGLE_MODEL,GOOGLE_CHAT_MIN_INTERVAL from auto_media_publisher.config.conf_base import GOOGLE_GEMINI_TEMPERATURE_MIN,GOOGLE_GEMINI_TEMPERATURE_MAX,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 class GeminiDocxRewriter: def __init__(self,rewrite_style="default" ): genai.configure(api_key=GOOGLE_API_KEY) self.temperature = round(random.uniform(GOOGLE_GEMINI_TEMPERATURE_MIN, GOOGLE_GEMINI_TEMPERATURE_MAX), 2) self.model = genai.GenerativeModel(GOOGLE_MODEL, generation_config=genai.types.GenerationConfig( temperature=self.temperature, top_p=1.0 ) ) if rewrite_style =="random": self.rewrite_style = random.choice(list(GOOGLE_REWRITE_STYLE_MAP.keys())) else: self.rewrite_style = rewrite_style 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() #@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,using temperature:{self.temperature}. 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) summary_info = self.summarize_title_intro_keywords() print(summary_info) self.text2imager.generate_image(summary_info["summary"],output_path.replace(".docx",".jpg")) summary_info["cover_image"] = output_path.replace(".docx",".jpg") # 3. 写入磁盘 with open(output_path.replace(".docx",".json"), "w", encoding="utf-8") as f: json.dump(summary_info, f, ensure_ascii=False, indent=4) return summary_info 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() for block in blocks: if block["type"]=="text": for line in block["content"].split("\n"): doc.add_paragraph(line.strip()) 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 doc.save(output_path) def summarize_title_intro_keywords(self) -> dict: title = self.safe_send_message(GEMINI_PROMPT_TITLE).text.strip() summary = self.safe_send_message(GEMINI_PROMPT_DESCRIPTION).text.strip() keywords = self.safe_send_message(GEMINI_PROMPT_TAG).text.strip() return { "title": title, "summary": summary, "keywords": re.split(r'[、,\s]+', keywords.strip("。")) }