Files
yu-zhi-ran/scripts/writer.py
T

325 lines
13 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
#!/usr/bin/env python3
"""
撰写阶段:基于大纲和选题生成完整文章(三平台版本)
"""
import json, datetime, logging, sys, re, subprocess
from pathlib import Path
from typing import Dict, List
import base64
from io import BytesIO
PROJECT_ROOT = Path(__file__).parent.parent
# 添加项目根和 backend 路径,以导入 app.core.llm_client
sys.path.insert(0, str(PROJECT_ROOT))
sys.path.insert(0, str(PROJECT_ROOT / 'platform' / 'backend'))
# 导入 LLM 客户端(NVIDIA
try:
from app.core.qnaigc_client import expand_content_with_llm # type: ignore
HAVE_LLM = True
except ImportError as e:
logging.warning(f"LLM client unavailable: {e}")
HAVE_LLM = False
PROJECT_ROOT = Path(__file__).parent.parent
sys.path.insert(0, str(PROJECT_ROOT))
# 导入数据库模型
from app.database import SessionLocal
from app.models import Topic
DATA_DIR = PROJECT_ROOT / "automation" / "data"
TOPICS_FILE = DATA_DIR / "sustainability_topics.json"
OUTLINE_DIR = DATA_DIR / "outlines"
RELEASE_DIR = DATA_DIR / "releases"
TEMPLATES_DIR = PROJECT_ROOT / "automation" / "templates"
LOGS_DIR = PROJECT_ROOT / "automation" / "logs"
TODAY = datetime.datetime.now().strftime("%Y-%m-%d")
logging.basicConfig(
level=logging.INFO,
format='%(asctime)s - %(levelname)s - %(message)s',
handlers=[
logging.FileHandler(LOGS_DIR / f"writer_{TODAY}.log"),
logging.StreamHandler()
]
)
logger = logging.getLogger(__name__)
class Writer:
def __init__(self, topic_id: str):
self.topic_id = topic_id
self.topic = self._load_topic()
outline_file = OUTLINE_DIR / TODAY / f"{topic_id}_outline.md"
if not outline_file.exists():
raise FileNotFoundError(f"Outline not found: {outline_file}")
self.outline_content = outline_file.read_text(encoding='utf-8')
self.release_dir = RELEASE_DIR / TODAY
self.release_dir.mkdir(parents=True, exist_ok=True)
# 加载研究笔记(作为 LLM 上下文)
research_file = DATA_DIR / "research" / TODAY / f"{topic_id}_research.md"
self.research_notes = research_file.read_text(encoding='utf-8') if research_file.exists() else ""
def _load_topic(self) -> Dict:
topics = json.loads(TOPICS_FILE.read_text(encoding='utf-8'))
for t in topics:
if t['id'] == self.topic_id:
return t
raise ValueError(f"Topic {self.topic_id} not found")
def _clean_title(self, title: str) -> str:
"""去除标题中的指导性文字(如字数说明、MVP标记等)"""
import re
# 去掉括号中的说明:约200字、约300字、MVP、试行等
title = re.sub(r'[(]约\s*\d+字[)]', '', title)
title = re.sub(r'[(]MVP[)]', '', title)
title = re.sub(r'[(][^)]*?[)]', '', title) # 保守移除任意括号内容(可能误伤,但大纲通常不包含重要括号信息)
return title.strip()
def _parse_outline_sections(self) -> List[Dict]:
"""将大纲 Markdown 解析为结构化列表,保留层级和内容"""
sections = []
current = None
for line in self.outline_content.splitlines():
if line.startswith("# "):
if current:
sections.append(current)
current = {"level": 1, "title": line[2:].strip(), "content": ""}
elif line.startswith("## "):
if current:
sections.append(current)
current = {"level": 2, "title": line[3:].strip(), "content": ""}
elif line.startswith("### "):
if current:
sections.append(current)
current = {"level": 3, "title": line[4:].strip(), "content": ""}
else:
if current and line.strip():
current['content'] = current.get('content', '') + line + "\n"
if current:
sections.append(current)
return sections
def _expand_section(self, section: Dict) -> str:
"""将大纲中的简短描述扩展为完整段落"""
content = section.get('content', '').strip()
# 如果有足够内容(>200字),直接返回
if len(content) > 200:
return content
# 如果内容极少,需要 LLM 扩写
if HAVE_LLM and len(content) < 150:
logger.info(f"使用 LLM 扩写章节: {section['title']}")
try:
expanded = expand_content_with_llm(
topic=self.topic,
section_title=section['title'],
section_content=content,
context=self.research_notes
)
if expanded and len(expanded.strip()) > len(content):
return expanded.strip()
else:
logger.warning("LLM 扩写结果为空或过短,使用占位")
raise ValueError("Empty expansion")
except Exception as e:
logger.warning(f"LLM 扩写失败: {e},使用占位内容")
# 返回占位内容,保持流程继续
return f"{content}\n\n(本段内容需要人工补充:当前模型调用失败或未配置)"
# 否则返回原内容
return content
def generate_full_markdown(self) -> str:
"""根据大纲生成完整 Markdown 正文(不用原标题,全部由 LLM 扩写生成)"""
sections = self._parse_outline_sections()
parts = []
# 只保留 LLM 扩写的内容,不添加任何原始标题标记
for sec in sections:
# 如果内容极短,LLM 扩写后返回的完整段落中可能包含标题,我们不过滤
if sec.get('content'):
expanded = self._expand_section(sec)
parts.append(expanded + "\n\n")
full_md = "\n".join(parts).strip()
# 添加文末声明
full_md += f"\n<p>(本文由宇之然AI助手生成,数据来源可靠,内容经合规审查)</p>\n"
full_md += f"<p><em>生成时间:{TODAY}</em></p>\n"
return full_md
def generate_platform_html(self, markdown: str, platform: str) -> str:
"""将 Markdown 转换为平台 HTML(基于模板)"""
title = self.topic['title']
# 加载模板
tpl_path = TEMPLATES_DIR / f"{platform}.html"
if tpl_path.exists():
template = tpl_path.read_text(encoding='utf-8')
else:
template = "<!DOCTYPE html><html><head><meta charset='UTF-8'><title>{{TITLE}}</title></head><body><h1>{{TITLE}}</h1><!-- CONTENT --></body></html>"
# 替换变量
html = template.replace("{{TITLE}}", title).replace("{{DATE}}", TODAY)
# 注入内容 (简单处理:markdown 转 HTML 可以用 marked.js 或 simple转换,这里暂时用 <pre> 包裹或简单段落化)
# 为了快速展示,我们将 markdown 的段落转换为 <p> 标签
# 实际中建议使用 markdown 库(如 python-markdown)转换
html_content = self._markdown_to_html(markdown)
html = html.replace("<!-- CONTENT -->", html_content)
# 平台特定标签补充
if platform == "zhihu":
tags = '<div class="tags">#科技 #职场</div>'
html = html.replace("<!-- TAGS -->", tags)
elif platform == "xiaohongshu":
hashtags = '<div class="hashtags">#AI #可持续 #生活方式</div>'
html = html.replace("<!-- HASHTAGS -->", hashtags)
elif platform == "wechat":
# 微信公众号可能还需要摘要等,模板已处理
pass
if platform == "xiaohongshu":
html = self._fill_image_placeholders(html, platform, title)
return html
def _markdown_to_html(self, md: str) -> str:
"""极简 markdown 转换(仅本场景使用)"""
lines = md.split('\n')
html_parts = []
for line in lines:
if line.startswith('# '):
html_parts.append(f"<h1>{line[2:]}</h1>")
elif line.startswith('## '):
html_parts.append(f"<h2>{line[3:]}</h2>")
elif line.startswith('### '):
html_parts.append(f"<h3>{line[4:]}</h3>")
elif line.strip().startswith('- '):
html_parts.append(f"<li>{line[2:]}</li>")
elif re.match(r'^\d+\. ', line):
content = re.sub(r'^\d+\. ', '', line)
html_parts.append(f"<li>{content}</li>")
elif line.strip():
html_parts.append(f"<p>{line}</p>")
else:
html_parts.append("") # 空行
return "\n".join(html_parts)
def save_html(self, html: str, platform: str) -> Path:
out_dir = self.release_dir / platform
out_dir.mkdir(parents=True, exist_ok=True)
filename = f"{platform}_{self.topic_id}_{platform}.html"
out_path = out_dir / filename
out_path.write_text(html, encoding='utf-8')
logger.info(f"HTML 生成: {out_path}")
return out_path
def mark_draft(self):
"""标记选题为「待审查」,同时更新数据库"""
"""标记选题为「待发布」,同时更新数据库"""
# 更新 JSON 文件
with open(TOPICS_FILE, 'r', encoding='utf-8') as f:
topics = json.load(f)
updated = False
for t in topics:
if t.get('id') == self.topic_id:
t['status'] = '待审查'
updated = True
break
if updated:
with open(TOPICS_FILE, 'w', encoding='utf-8') as f:
json.dump(topics, f, ensure_ascii=False, indent=2)
# 更新数据库
db = SessionLocal()
try:
topic_db = db.query(Topic).filter(Topic.id == self.topic_id).first()
if topic_db:
topic_db.status = '待审查'
db.commit()
logger.info(f"选题 {self.topic_id} 状态已更新为 draft(数据库)")
else:
logger.warning(f"数据库中未找到选题 {self.topic_id}")
except Exception as e:
logger.error(f"更新数据库失败: {e}")
db.rollback()
finally:
db.close()
logger.info(f"选题 {self.topic_id} 状态更新为「待审查」(JSON")
"""标记选题为「待发布」"""
with open(TOPICS_FILE, 'r', encoding='utf-8') as f:
topics = json.load(f)
for t in topics:
if t.get('id') == self.topic_id:
t['status'] = '待审查'
# ready_at 留空,待合规审核通过后设置
break
with open(TOPICS_FILE, 'w', encoding='utf-8') as f:
json.dump(topics, f, ensure_ascii=False, indent=2)
logger.info(f"选题 {self.topic_id} 状态更新为「待审查」")
def run(self):
logger.info("开始撰写阶段")
markdown = self.generate_full_markdown()
results = {}
for platform in ["zhihu", "wechat", "xiaohongshu"]:
html = self.generate_platform_html(markdown, platform)
results[platform] = str(self.save_html(html, platform))
self.mark_draft()
logger.info(f"撰写完成,状态改为 draft,待合规审核")
return {"ok": True, "files": results}
def _image_to_data_url(self, img_path: Path, fmt: str = None) -> str:
data = img_path.read_bytes()
if fmt is None:
fmt = img_path.suffix.lstrip('.').lower()
b64 = base64.b64encode(data).decode('ascii')
return f"data:image/{fmt};base64,{b64}"
def _generate_and_inline_images(self, platform: str, title: str) -> dict:
from scripts.image_generator import ImageGenerator
gen = ImageGenerator()
files = gen.generate_all_placeholders(title, platform)
mapping = {}
cover = files.get('cover')
if cover and cover.exists():
mapping['main-image-src'] = self._image_to_data_url(cover)
thumbs = []
for k, p in files.items():
if k != 'cover' and p.exists():
thumbs.append(self._image_to_data_url(p))
mapping['thumbnail-srcs'] = thumbs
return mapping
def _fill_image_placeholders(self, html: str, platform: str, title: str) -> str:
if platform != 'xiaohongshu':
return html
mapping = self._generate_and_inline_images(platform, title)
# Replace main image placeholder
main_ph = '<img src="" alt="封面图" class="main-image">'
if 'main-image-src' in mapping:
new_main = f'<img src="{mapping["main-image-src"]}" alt="封面图" class="main-image">'
html = html.replace(main_ph, new_main)
# Replace thumbnail placeholders (6)
thumbs = mapping.get('thumbnail-srcs', [])
for idx, src in enumerate(thumbs[:6], start=1):
ph = f'<img src="" alt="{idx}" class="thumbnail">'
new_thumb = f'<img src="{src}" alt="{idx}" class="thumbnail">'
html = html.replace(ph, new_thumb)
return html
def main():
import argparse
parser = argparse.ArgumentParser()
parser.add_argument('--topic-id', required=True, help='选题ID')
args = parser.parse_args()
w = Writer(args.topic_id)
result = w.run()
print(json.dumps(result, ensure_ascii=False))
sys.exit(0 if result['ok'] else 1)
if __name__ == "__main__":
main()