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 import google.generativeai as genai from auto_media_publisher.config.conf_base import GOOGLE_API_KEY,GOOGLE_MODEL,GOOGLE_CHAT_MIN_INTERVAL,DEFAULT_PREFIX_MAP,DEFAULT_SURFIX_MAP from auto_media_publisher.config.conf_base import GOOGLE_GEMINI_TEMPERATURE_MIN,GOOGLE_GEMINI_TEMPERATURE_MAX,GOOGLE_GEMINI_TOPN_MIN,GOOGLE_GEMINI_TOPN_MAX 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 class DocxGeneratorGemini: def __init__(self,docx_author="我",complex=False): genai.configure(api_key=GOOGLE_API_KEY) self.temperature = round(random.uniform(GOOGLE_GEMINI_TEMPERATURE_MIN, GOOGLE_GEMINI_TEMPERATURE_MAX), 2) self.top_n = round(random.uniform(GOOGLE_GEMINI_TOPN_MIN, GOOGLE_GEMINI_TOPN_MAX), 2) self.model = genai.GenerativeModel(GOOGLE_MODEL, generation_config=genai.types.GenerationConfig( temperature=self.temperature, top_p=self.top_n, top_k=random.choice([40, 50, 80, 100]) ) ) 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 rewriting docx using gemini,using temperature:{self.temperature}. 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("。")) }