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_with_trans(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("。"))
|
}
|
