261 lines
11 KiB
Python
261 lines
11 KiB
Python
|
#!/usr/bin/env python3
|
|||
|
|
# -*- coding: utf-8 -*-
|
|||
|
|
"""
|
|||
|
|
测试选题生成和文章生成的脚本
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
import os
|
|||
|
|
import time
|
|||
|
|
import argparse
|
|||
|
|
import json
|
|||
|
|
import logging
|
|||
|
|
import sys
|
|||
|
|
from datetime import datetime
|
|||
|
|
|
|||
|
|
from core.ai_agent import AI_Agent
|
|||
|
|
from utils.prompt_manager import PromptManager
|
|||
|
|
from utils.tweet_generator import run_topic_generation_pipeline, generate_content_for_topic
|
|||
|
|
from utils.output_handler import FileSystemOutputHandler
|
|||
|
|
|
|||
|
|
def load_config(config_path="topic_content_config.json"):
|
|||
|
|
"""从JSON文件加载配置"""
|
|||
|
|
if not os.path.exists(config_path):
|
|||
|
|
print(f"错误：配置文件 '{config_path}' 未找到。")
|
|||
|
|
sys.exit(1)
|
|||
|
|
try:
|
|||
|
|
with open(config_path, 'r', encoding='utf-8') as f:
|
|||
|
|
config = json.load(f)
|
|||
|
|
|
|||
|
|
# 基本验证
|
|||
|
|
required_keys = ["api_url", "model", "api_key", "resource_dir", "output_dir",
|
|||
|
|
"num", "variants", "topic_system_prompt", "topic_user_prompt",
|
|||
|
|
"content_system_prompt"]
|
|||
|
|
|
|||
|
|
if not all(key in config for key in required_keys):
|
|||
|
|
missing_keys = [key for key in required_keys if key not in config]
|
|||
|
|
print(f"错误：配置文件 '{config_path}' 缺少必需的键：{missing_keys}")
|
|||
|
|
sys.exit(1)
|
|||
|
|
|
|||
|
|
# 验证prompts_dir或prompts_config至少有一个存在
|
|||
|
|
if not ("prompts_dir" in config or "prompts_config" in config):
|
|||
|
|
print(f"错误：配置文件 '{config_path}' 必须包含 'prompts_dir' 或 'prompts_config'")
|
|||
|
|
sys.exit(1)
|
|||
|
|
|
|||
|
|
return config
|
|||
|
|
except json.JSONDecodeError:
|
|||
|
|
print(f"错误：无法从 '{config_path}' 解码JSON。请检查文件格式。")
|
|||
|
|
sys.exit(1)
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"从 '{config_path}' 加载配置时出错：{e}")
|
|||
|
|
sys.exit(1)
|
|||
|
|
|
|||
|
|
def main():
|
|||
|
|
# 设置日志记录
|
|||
|
|
logging.basicConfig(
|
|||
|
|
level=logging.INFO,
|
|||
|
|
format='%(asctime)s - %(levelname)s - [%(filename)s:%(lineno)d] - %(message)s',
|
|||
|
|
datefmt='%Y-%m-%d %H:%M:%S'
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# 解析命令行参数
|
|||
|
|
parser = argparse.ArgumentParser(description="测试选题和文章生成")
|
|||
|
|
parser.add_argument(
|
|||
|
|
"--config",
|
|||
|
|
type=str,
|
|||
|
|
default="topic_content_config.json",
|
|||
|
|
help="配置文件路径（例如，topic_content_config.json）"
|
|||
|
|
)
|
|||
|
|
parser.add_argument(
|
|||
|
|
"--run_id",
|
|||
|
|
type=str,
|
|||
|
|
default=None,
|
|||
|
|
help="可选的指定运行ID（例如，'test_run_01'）。如果未提供，将生成基于时间戳的ID。"
|
|||
|
|
)
|
|||
|
|
parser.add_argument(
|
|||
|
|
"--topics_file",
|
|||
|
|
type=str,
|
|||
|
|
default=None,
|
|||
|
|
help="可选的预生成选题JSON文件路径。如果提供，则跳过选题生成。"
|
|||
|
|
)
|
|||
|
|
parser.add_argument(
|
|||
|
|
"--debug",
|
|||
|
|
action='store_true',
|
|||
|
|
help="启用调试级别日志记录。"
|
|||
|
|
)
|
|||
|
|
args = parser.parse_args()
|
|||
|
|
|
|||
|
|
# 调整日志级别（如果启用了调试）
|
|||
|
|
if args.debug:
|
|||
|
|
logging.getLogger().setLevel(logging.DEBUG)
|
|||
|
|
logging.info("已启用调试日志记录。")
|
|||
|
|
|
|||
|
|
logging.info("启动选题和文章生成测试脚本...")
|
|||
|
|
logging.info(f"使用配置文件：{args.config}")
|
|||
|
|
if args.run_id:
|
|||
|
|
logging.info(f"使用指定的run_id：{args.run_id}")
|
|||
|
|
if args.topics_file:
|
|||
|
|
logging.info(f"使用现有选题文件：{args.topics_file}")
|
|||
|
|
|
|||
|
|
# 加载配置
|
|||
|
|
config = load_config(args.config)
|
|||
|
|
if config is None:
|
|||
|
|
logging.critical("无法加载配置。退出。")
|
|||
|
|
sys.exit(1)
|
|||
|
|
|
|||
|
|
# 初始化输出处理器
|
|||
|
|
output_handler = FileSystemOutputHandler(config.get("output_dir", "result"))
|
|||
|
|
logging.info(f"使用输出处理器：{output_handler.__class__.__name__}")
|
|||
|
|
|
|||
|
|
run_id = args.run_id
|
|||
|
|
topics_list = None
|
|||
|
|
system_prompt = None
|
|||
|
|
user_prompt = None
|
|||
|
|
pipeline_start_time = time.time()
|
|||
|
|
|
|||
|
|
# 步骤1：选题生成（或加载现有选题）
|
|||
|
|
if args.topics_file:
|
|||
|
|
from core.topic_parser import TopicParser
|
|||
|
|
logging.info(f"跳过选题生成（步骤1）- 从以下位置加载选题：{args.topics_file}")
|
|||
|
|
topics_list = TopicParser.load_topics_from_json(args.topics_file)
|
|||
|
|
if topics_list:
|
|||
|
|
# 如果未提供run_id，尝试从文件名推断
|
|||
|
|
if not run_id:
|
|||
|
|
try:
|
|||
|
|
base = os.path.basename(args.topics_file)
|
|||
|
|
# 假设格式为"tweet_topic_{run_id}.json"或"tweet_topic.json"
|
|||
|
|
if base.startswith("tweet_topic_") and base.endswith(".json"):
|
|||
|
|
run_id = base[len("tweet_topic_"): -len(".json")]
|
|||
|
|
logging.info(f"从选题文件名推断的run_id：{run_id}")
|
|||
|
|
elif base == "tweet_topic.json":
|
|||
|
|
logging.warning(f"从默认文件名'{base}'加载选题。未推断run_id。")
|
|||
|
|
else:
|
|||
|
|
logging.warning(f"无法从选题文件名推断run_id：{base}")
|
|||
|
|
except Exception as e:
|
|||
|
|
logging.warning(f"尝试从选题文件名推断run_id时出错：{e}")
|
|||
|
|
|
|||
|
|
# 如果尝试推断后run_id仍为None，则生成一个
|
|||
|
|
if run_id is None:
|
|||
|
|
run_id = datetime.now().strftime("%Y-%m-%d_%H-%M-%S_loaded")
|
|||
|
|
logging.info(f"为加载的选题生成的run_id：{run_id}")
|
|||
|
|
|
|||
|
|
# 加载文件时缺少提示词
|
|||
|
|
system_prompt = ""
|
|||
|
|
user_prompt = ""
|
|||
|
|
logging.info(f"成功加载{len(topics_list)}个选题，run_id：{run_id}。提示词不可用。")
|
|||
|
|
else:
|
|||
|
|
logging.error(f"无法从{args.topics_file}加载选题。无法继续。")
|
|||
|
|
sys.exit(1)
|
|||
|
|
else:
|
|||
|
|
logging.info("执行选题生成（步骤1）...")
|
|||
|
|
step1_start = time.time()
|
|||
|
|
# 调用更新后的函数，接收原始数据
|
|||
|
|
run_id, topics_list, system_prompt, user_prompt = run_topic_generation_pipeline(config, args.run_id)
|
|||
|
|
step1_end = time.time()
|
|||
|
|
if run_id is not None and topics_list is not None:
|
|||
|
|
logging.info(f"步骤1成功完成，耗时{step1_end - step1_start:.2f}秒。运行ID：{run_id}")
|
|||
|
|
# 使用输出处理器保存结果
|
|||
|
|
output_handler.handle_topic_results(run_id, topics_list, system_prompt, user_prompt)
|
|||
|
|
else:
|
|||
|
|
logging.critical("选题生成（步骤1）失败。退出。")
|
|||
|
|
sys.exit(1)
|
|||
|
|
|
|||
|
|
# 步骤2：内容生成
|
|||
|
|
if run_id is not None and topics_list is not None:
|
|||
|
|
logging.info("执行内容生成（步骤2）...")
|
|||
|
|
step2_start = time.time()
|
|||
|
|
|
|||
|
|
# 创建PromptManager实例
|
|||
|
|
try:
|
|||
|
|
prompt_manager = PromptManager(
|
|||
|
|
topic_system_prompt_path=config.get("topic_system_prompt"),
|
|||
|
|
topic_user_prompt_path=config.get("topic_user_prompt"),
|
|||
|
|
content_system_prompt_path=config.get("content_system_prompt"),
|
|||
|
|
prompts_config=config.get("prompts_config"),
|
|||
|
|
prompts_dir=config.get("prompts_dir"),
|
|||
|
|
resource_dir_config=config.get("resource_dir", []),
|
|||
|
|
topic_gen_num=config.get("num", 1),
|
|||
|
|
topic_gen_date=config.get("date", "")
|
|||
|
|
)
|
|||
|
|
logging.info("已为步骤2创建PromptManager实例。")
|
|||
|
|
except KeyError as e:
|
|||
|
|
logging.error(f"创建PromptManager时配置错误：缺少键'{e}'。无法继续内容生成。")
|
|||
|
|
sys.exit(1)
|
|||
|
|
|
|||
|
|
# 初始化AI Agent
|
|||
|
|
ai_agent = None
|
|||
|
|
content_success = False
|
|||
|
|
try:
|
|||
|
|
request_timeout = config.get("request_timeout", 30)
|
|||
|
|
max_retries = config.get("max_retries", 3)
|
|||
|
|
stream_chunk_timeout = config.get("stream_chunk_timeout", 60)
|
|||
|
|
ai_agent = AI_Agent(
|
|||
|
|
config["api_url"],
|
|||
|
|
config["model"],
|
|||
|
|
config["api_key"],
|
|||
|
|
timeout=request_timeout,
|
|||
|
|
max_retries=max_retries,
|
|||
|
|
stream_chunk_timeout=stream_chunk_timeout
|
|||
|
|
)
|
|||
|
|
logging.info("已初始化用于内容生成的AI Agent。")
|
|||
|
|
|
|||
|
|
# 遍历选题
|
|||
|
|
for i, topic_item in enumerate(topics_list):
|
|||
|
|
topic_index = topic_item.get('index', i + 1)
|
|||
|
|
logging.info(f"--- 处理选题 {topic_index}/{len(topics_list)}: {topic_item.get('object', 'N/A')} ---")
|
|||
|
|
|
|||
|
|
# 读取内容生成所需的参数
|
|||
|
|
content_variants = config.get("variants", 1)
|
|||
|
|
content_temp = config.get("content_temperature", 0.3)
|
|||
|
|
content_top_p = config.get("content_top_p", 0.4)
|
|||
|
|
content_presence_penalty = config.get("content_presence_penalty", 1.5)
|
|||
|
|
|
|||
|
|
# 调用generate_content_for_topic
|
|||
|
|
topic_success = generate_content_for_topic(
|
|||
|
|
ai_agent,
|
|||
|
|
prompt_manager,
|
|||
|
|
topic_item,
|
|||
|
|
run_id,
|
|||
|
|
topic_index,
|
|||
|
|
output_handler,
|
|||
|
|
variants=content_variants,
|
|||
|
|
temperature=content_temp,
|
|||
|
|
top_p=content_top_p,
|
|||
|
|
presence_penalty=content_presence_penalty
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
if topic_success:
|
|||
|
|
logging.info(f"选题{topic_index}的内容生成成功。")
|
|||
|
|
content_success = True
|
|||
|
|
else:
|
|||
|
|
logging.warning(f"选题{topic_index}的内容生成失败或未产生有效结果。")
|
|||
|
|
logging.info(f"--- 完成选题 {topic_index} ---")
|
|||
|
|
|
|||
|
|
except KeyError as e:
|
|||
|
|
logging.error(f"内容生成过程中的配置错误：缺少键'{e}'")
|
|||
|
|
traceback.print_exc()
|
|||
|
|
except Exception as e:
|
|||
|
|
logging.exception("内容生成过程中发生意外错误:")
|
|||
|
|
finally:
|
|||
|
|
# 确保AI agent已关闭
|
|||
|
|
if ai_agent:
|
|||
|
|
logging.info("关闭内容生成AI Agent...")
|
|||
|
|
ai_agent.close()
|
|||
|
|
|
|||
|
|
step2_end = time.time()
|
|||
|
|
if content_success:
|
|||
|
|
logging.info(f"步骤2完成，耗时{step2_end - step2_start:.2f}秒。")
|
|||
|
|
else:
|
|||
|
|
logging.warning("步骤2完成，但可能遇到错误或未生成输出。")
|
|||
|
|
else:
|
|||
|
|
logging.error("无法进行步骤2：步骤1的run_id或topics_list无效。")
|
|||
|
|
|
|||
|
|
# 最终化输出
|
|||
|
|
if run_id:
|
|||
|
|
output_handler.finalize(run_id)
|
|||
|
|
|
|||
|
|
pipeline_end_time = time.time()
|
|||
|
|
logging.info(f"流程完成。总执行时间：{pipeline_end_time - pipeline_start_time:.2f}秒。")
|
|||
|
|
logging.info(f"运行ID'{run_id}'的结果位于：{os.path.join(config.get('output_dir', 'result'), run_id)}")
|
|||
|
|
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
main()
|