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
|
from io import BytesIO
|
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 GEMINI_PROMPT_TITLE,GEMINI_PROMPT_DESCRIPTION,GEMINI_PROMPT_TAG,GOOGLE_STORY_GENERATE_PROMPT,GOOGLE_STORY_GENERATE_PROMPT_COMPLEX
|
from auto_media_publisher.config.conf_docx_generate import STORY_PROMPT_PARAMETERS
|
from auto_media_publisher.config.conf_docx_generate_complex import STORY_PROMPT_PARAMETERS_COMPLEX
|
|
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 DocxGeneratorGemini:
|
def __init__(self,docx_author="小白",complex=False):
|
|
self.model = init_gemini_model(GOOGLE_GEMINI_PROXY)
|
|
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 = {}
|
|
if complex :
|
self.prompt = self.get_story_prompt_complex()
|
else:
|
self.prompt = self.get_story_prompt()
|
|
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 get_story_prompt(self) -> str:
|
who = random.choice(STORY_PROMPT_PARAMETERS["who"])
|
what = random.choice(STORY_PROMPT_PARAMETERS["what"])
|
style = "，".join(random.sample(STORY_PROMPT_PARAMETERS["style"], 2))
|
level = random.choice(STORY_PROMPT_PARAMETERS["level"])
|
|
return GOOGLE_STORY_GENERATE_PROMPT.format(who=who,what=what,style=style,level=level)
|
|
def get_story_prompt_complex(self) -> str:
|
|
params = {key: random.choice(values) for key, values in STORY_PROMPT_PARAMETERS_COMPLEX.items()}
|
|
return GOOGLE_STORY_GENERATE_PROMPT_COMPLEX.format(**params)
|
|
|
def generate(self, output_path: str):
|
self.logger.info(f"start to Generating docx using gemini, with prompt {self.prompt}.")
|
self.chat = self.model.start_chat(history=[])
|
response = self.safe_send_message(self.prompt)
|
story_text = response.text.strip()
|
self.logger.info(f"[ google ] gemini story end...")
|
paragraphs = [p.strip() for p in story_text.split('\n') if p.strip()]
|
|
doc = Document()
|
# 设置默认样式
|
style = doc.styles['Normal']
|
font = style.font
|
font.name = '楷体' # 西文字体设置为楷体
|
font.size = Pt(10.5) # 小四字体
|
style._element.rPr.rFonts.set(qn('w:eastAsia'), '楷体') # 中文字体设定
|
|
self._add_content_to_doc(doc,"text",f"{self.prefix_text}\n")
|
|
for para in paragraphs :
|
self._add_content_to_doc(doc,"text",para)
|
if random.random() > 0.5:
|
try:
|
img_stream = self.text2imager.generate_image_to_stream(para)
|
self._add_content_to_doc(doc,"image",img_stream)
|
except:
|
pass
|
|
self._add_content_to_doc(doc,"text",f"\n\n{self.surfix_text}")
|
|
doc.save(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 _add_content_to_doc(self,doc,type,block_content):
|
if 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(5) # 段前0磅
|
para.paragraph_format.space_after = Pt(5) # 段后0磅
|
elif type == "image":
|
paragraph = doc.add_paragraph()
|
run = paragraph.add_run()
|
run.add_picture(block_content, width=Inches(5.5))
|
paragraph.paragraph_format.space_before = Pt(12) # 段前 12pt
|
paragraph.paragraph_format.space_after = Pt(12) # 段后 12pt
|
paragraph.alignment = WD_ALIGN_PARAGRAPH.CENTER
|
|
return doc
|
|
|
|
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("。"))
|
}
|
