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