383 lines
15 KiB
Python
383 lines
15 KiB
Python
#!/usr/bin/env python
|
||
# -*- coding: utf-8 -*-
|
||
|
||
import os
|
||
import time
|
||
import random
|
||
import argparse
|
||
import json
|
||
from datetime import datetime
|
||
import sys
|
||
import traceback
|
||
sys.path.append('/root/autodl-tmp')
|
||
# 从本地模块导入
|
||
from TravelContentCreator.core.ai_agent import AI_Agent
|
||
from TravelContentCreator.core.topic_parser import TopicParser
|
||
# ResourceLoader is now used implicitly via PromptManager
|
||
# from TravelContentCreator.utils.resource_loader import ResourceLoader
|
||
from TravelContentCreator.utils.prompt_manager import PromptManager # Import PromptManager
|
||
|
||
class tweetTopic:
|
||
def __init__(self, index, date, logic, object, product, product_logic, style, style_logic, target_audience, target_audience_logic):
|
||
self.index = index
|
||
self.date = date
|
||
self.logic = logic
|
||
self.object = object
|
||
self.product = product
|
||
self.product_logic = product_logic
|
||
self.style = style
|
||
self.style_logic = style_logic
|
||
self.target_audience = target_audience
|
||
self.target_audience_logic = target_audience_logic
|
||
|
||
class tweetTopicRecord:
|
||
def __init__(self, topics_list, system_prompt, user_prompt, output_dir, run_id):
|
||
self.topics_list = topics_list
|
||
self.system_prompt = system_prompt
|
||
self.user_prompt = user_prompt
|
||
self.output_dir = output_dir
|
||
self.run_id = run_id
|
||
|
||
def save_topics(self, path):
|
||
try:
|
||
with open(path, "w", encoding="utf-8") as f:
|
||
json.dump(self.topics_list, f, ensure_ascii=False, indent=4)
|
||
except Exception as e:
|
||
print(f"保存选题失败: {e}")
|
||
return False
|
||
return True
|
||
|
||
def save_prompt(self, path):
|
||
try:
|
||
with open(path, "w", encoding="utf-8") as f:
|
||
f.write(self.system_prompt + "\n")
|
||
f.write(self.user_prompt + "\n")
|
||
f.write(self.output_dir + "\n")
|
||
f.write(self.run_id + "\n")
|
||
except Exception as e:
|
||
print(f"保存提示词失败: {e}")
|
||
return False
|
||
return True
|
||
|
||
class tweetContent:
|
||
def __init__(self, result, prompt, output_dir, run_id, article_index, variant_index):
|
||
self.result = result
|
||
self.prompt = prompt
|
||
self.output_dir = output_dir
|
||
self.run_id = run_id
|
||
self.article_index = article_index
|
||
self.variant_index = variant_index
|
||
|
||
self.title, self.content = self.split_content(result)
|
||
self.json_file = self.gen_result_json()
|
||
def split_content(self, result):
|
||
## remove <\think>
|
||
result = result.split("</think>")[1]
|
||
## get tile
|
||
title = result.split("title>")[1].split("</title>")[0]
|
||
## get content
|
||
content = result.split("content>")[1].split("</content>")[0]
|
||
return title, content
|
||
|
||
def gen_result_json(self):
|
||
json_file = {
|
||
"title": self.title,
|
||
"content": self.content
|
||
}
|
||
return json_file
|
||
|
||
def save_content(self, json_path):
|
||
with open(json_path, "w", encoding="utf-8") as f:
|
||
json.dump(self.json_file, f, ensure_ascii=False, indent=4)
|
||
return json_path
|
||
|
||
def save_prompt(self, path):
|
||
with open(path, "w", encoding="utf-8") as f:
|
||
f.write(self.prompt + "\n")
|
||
return path
|
||
|
||
def get_content(self):
|
||
return self.content
|
||
|
||
def get_title(self):
|
||
return self.title
|
||
|
||
def get_json_file(self):
|
||
return self.json_file
|
||
|
||
|
||
def generate_topics(ai_agent, system_prompt, user_prompt, output_dir, temperature=0.2, top_p=0.5, presence_penalty=1.5):
|
||
"""生成选题列表"""
|
||
print("开始生成选题...")
|
||
|
||
# 记录开始时间
|
||
time_start = time.time()
|
||
|
||
# 生成选题
|
||
result, system_prompt, user_prompt, file_folder, file_name, tokens, time_cost = ai_agent.work(
|
||
system_prompt, user_prompt, "", temperature, top_p, presence_penalty
|
||
)
|
||
|
||
# 计算总耗时
|
||
time_end = time.time()
|
||
print(f"选题生成完成，耗时：{time_end - time_start}秒")
|
||
|
||
# 生成唯一ID
|
||
run_id = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
|
||
|
||
# 解析选题
|
||
result_list = TopicParser.parse_topics(result)
|
||
# success, json_path = TopicParser.save_topics(result_list, output_dir, run_id, result)
|
||
|
||
tweet_topic_record = tweetTopicRecord(result_list, system_prompt, user_prompt, output_dir, run_id)
|
||
|
||
return run_id, tweet_topic_record
|
||
|
||
|
||
def generate_single_content(ai_agent, system_prompt, user_prompt, item, output_dir, run_id,
|
||
article_index, variant_index, temperature=0.3, top_p=0.4, presence_penalty=1.5):
|
||
"""生成单篇文章内容. Requires prompts to be passed in."""
|
||
try:
|
||
# Prompts are now passed directly as arguments
|
||
# No longer build user_prompt here
|
||
# user_prompt = ResourceLoader.build_user_prompt(item, prompts_dir, resource_dir)
|
||
|
||
if not system_prompt or not user_prompt:
|
||
print("Error: System or User prompt is empty. Cannot generate content.")
|
||
return None, None
|
||
|
||
print(f"Using pre-constructed prompts. User prompt length: {len(user_prompt)}")
|
||
|
||
# 添加随机停顿，避免请求过于频繁
|
||
time.sleep(random.random() * 0.5 + 0.1)
|
||
|
||
# 生成文章
|
||
result, _, _, _, _, tokens, time_cost = ai_agent.work(
|
||
system_prompt, user_prompt, "", temperature, top_p, presence_penalty
|
||
)
|
||
|
||
print(f"生成完成，tokens: {tokens}, 耗时: {time_cost}s")
|
||
|
||
# 保存到单独文件
|
||
tweet_content = tweetContent(result, user_prompt, output_dir, run_id, article_index, variant_index)
|
||
result_dir = os.path.join(output_dir, f"{article_index}_{variant_index}")
|
||
os.makedirs(result_dir, exist_ok=True)
|
||
tweet_content.save_content(os.path.join(result_dir, "article.json"))
|
||
tweet_content.save_prompt(os.path.join(result_dir, "tweet_prompt.txt"))
|
||
|
||
return tweet_content, result
|
||
except Exception as e:
|
||
print(f"生成单篇文章时出错: {e}")
|
||
return None, None
|
||
|
||
def generate_content(ai_agent, system_prompt, topics, output_dir, run_id, prompts_dir, resource_dir,
|
||
variants=2, temperature=0.3, start_index=0, end_index=None):
|
||
"""根据选题生成内容"""
|
||
if not topics:
|
||
print("没有选题，无法生成内容")
|
||
return
|
||
|
||
# 确定处理范围
|
||
if end_index is None or end_index > len(topics):
|
||
end_index = len(topics)
|
||
|
||
topics_to_process = topics[start_index:end_index]
|
||
print(f"准备处理{len(topics_to_process)}个选题...")
|
||
|
||
# 创建汇总文件
|
||
# summary_file = ResourceLoader.create_summary_file(output_dir, run_id, len(topics_to_process))
|
||
|
||
# 处理每个选题
|
||
processed_results = []
|
||
for i, item in enumerate(topics_to_process):
|
||
print(f"处理第 {i+1}/{len(topics_to_process)} 篇文章")
|
||
|
||
# 为每个选题生成多个变体
|
||
for j in range(variants):
|
||
print(f"正在生成变体 {j+1}/{variants}")
|
||
|
||
# 调用单篇文章生成函数
|
||
tweet_content, result = generate_single_content(
|
||
ai_agent, system_prompt, item, output_dir, run_id, i+1, j+1, temperature
|
||
)
|
||
|
||
if tweet_content:
|
||
processed_results.append(tweet_content)
|
||
|
||
# # 更新汇总文件 (仅保存第一个变体到汇总文件)
|
||
# if j == 0:
|
||
# ResourceLoader.update_summary(summary_file, i+1, user_prompt, result)
|
||
|
||
print(f"完成{len(processed_results)}篇文章生成")
|
||
return processed_results
|
||
|
||
|
||
def prepare_topic_generation(
|
||
config # Pass the whole config dictionary now
|
||
# select_date, select_num,
|
||
# system_prompt_path, user_prompt_path,
|
||
# base_url="vllm", model_name="qwenQWQ", api_key="EMPTY",
|
||
# gen_prompts_path="/root/autodl-tmp/TravelContentCreator/genPrompts",
|
||
# resource_dir="/root/autodl-tmp/TravelContentCreator/resource",
|
||
# output_dir="/root/autodl-tmp/TravelContentCreator/result"
|
||
):
|
||
"""准备选题生成的环境和参数. Returns agent and prompts."""
|
||
|
||
# Initialize PromptManager
|
||
prompt_manager = PromptManager(config)
|
||
|
||
# Get prompts using PromptManager
|
||
system_prompt, user_prompt = prompt_manager.get_topic_prompts()
|
||
|
||
if not system_prompt or not user_prompt:
|
||
print("Error: Failed to get topic generation prompts.")
|
||
return None, None, None, None
|
||
|
||
# 创建AI Agent (still create agent here for the topic generation phase)
|
||
try:
|
||
print("Initializing AI Agent for topic generation...")
|
||
ai_agent = AI_Agent(config["api_url"], config["model"], config["api_key"])
|
||
except Exception as e:
|
||
print(f"Error initializing AI Agent for topic generation: {e}")
|
||
traceback.print_exc()
|
||
return None, None, None, None
|
||
|
||
# Removed prompt loading/building logic, now handled by PromptManager
|
||
|
||
# Return agent and the generated prompts
|
||
return ai_agent, system_prompt, user_prompt, config["output_dir"]
|
||
|
||
def run_topic_generation_pipeline(config):
|
||
"""Runs the complete topic generation pipeline based on the configuration."""
|
||
print("Step 1: Generating Topics...")
|
||
|
||
# Prepare necessary inputs and the AI agent for topic generation
|
||
ai_agent, system_prompt, user_prompt, base_output_dir = None, None, None, None
|
||
try:
|
||
# Pass the config directly to prepare_topic_generation
|
||
ai_agent, system_prompt, user_prompt, base_output_dir = prepare_topic_generation(config)
|
||
|
||
if not ai_agent or not system_prompt or not user_prompt:
|
||
raise ValueError("Failed to prepare topic generation (agent or prompts missing).")
|
||
|
||
except Exception as e:
|
||
print(f"Error during topic generation preparation: {e}")
|
||
traceback.print_exc()
|
||
return None, None
|
||
|
||
# Generate topics using the prepared agent and prompts
|
||
try:
|
||
run_id, tweet_topic_record = generate_topics(
|
||
ai_agent, system_prompt, user_prompt, config["output_dir"],
|
||
config.get("topic_temperature", 0.2),
|
||
config.get("topic_top_p", 0.5),
|
||
config.get("topic_max_tokens", 1.5) # Consider if max_tokens name is accurate here (was presence_penalty?)
|
||
)
|
||
except Exception as e:
|
||
print(f"Error during topic generation API call: {e}")
|
||
traceback.print_exc()
|
||
if ai_agent: ai_agent.close() # Ensure agent is closed on error
|
||
return None, None
|
||
|
||
# Ensure the AI agent is closed after generation
|
||
if ai_agent:
|
||
ai_agent.close()
|
||
|
||
# Process results
|
||
if not run_id or not tweet_topic_record:
|
||
print("Topic generation failed (no run_id or topics returned).")
|
||
return None, None
|
||
|
||
output_dir = os.path.join(config["output_dir"], run_id)
|
||
try:
|
||
os.makedirs(output_dir, exist_ok=True)
|
||
# Save topics and prompt details
|
||
save_topics_success = tweet_topic_record.save_topics(os.path.join(output_dir, "tweet_topic.json"))
|
||
save_prompt_success = tweet_topic_record.save_prompt(os.path.join(output_dir, "tweet_prompt.txt"))
|
||
if not save_topics_success or not save_prompt_success:
|
||
print("Warning: Failed to save topic generation results or prompts.")
|
||
# Continue but warn user
|
||
|
||
except Exception as e:
|
||
print(f"Error saving topic generation results: {e}")
|
||
traceback.print_exc()
|
||
# Return the generated data even if saving fails, but maybe warn more strongly?
|
||
# return run_id, tweet_topic_record # Decide if partial success is okay
|
||
return None, None # Or consider failure if saving is critical
|
||
|
||
print(f"Topics generated successfully. Run ID: {run_id}")
|
||
return run_id, tweet_topic_record
|
||
|
||
def main():
|
||
"""主函数入口"""
|
||
config_file = {
|
||
"date": "4月17日",
|
||
"num": 5,
|
||
"model": "qwenQWQ",
|
||
"api_url": "vllm",
|
||
"api_key": "EMPTY",
|
||
"topic_system_prompt": "/root/autodl-tmp/TravelContentCreator/SelectPrompt/systemPrompt.txt",
|
||
"topic_user_prompt": "/root/autodl-tmp/TravelContentCreator/SelectPrompt/userPrompt.txt",
|
||
"content_system_prompt": "/root/autodl-tmp/TravelContentCreator/genPrompts/systemPrompt.txt",
|
||
"resource_dir": [{
|
||
"type": "Object",
|
||
"num": 4,
|
||
"file_path": ["/root/autodl-tmp/TravelContentCreator/resource/Object/景点信息-尚书第.txt",
|
||
"/root/autodl-tmp/TravelContentCreator/resource/Object/景点信息-明清园.txt",
|
||
"/root/autodl-tmp/TravelContentCreator/resource/Object/景点信息-泰宁古城.txt",
|
||
"/root/autodl-tmp/TravelContentCreator/resource/Object/景点信息-甘露寺.txt"
|
||
]},
|
||
{
|
||
"type": "Product",
|
||
"num": 0,
|
||
"file_path": []
|
||
}
|
||
],
|
||
"prompts_dir": "/root/autodl-tmp/TravelContentCreator/genPrompts",
|
||
"output_dir": "/root/autodl-tmp/TravelContentCreator/result",
|
||
"variants": 2,
|
||
"topic_temperature": 0.2,
|
||
"content_temperature": 0.3
|
||
}
|
||
|
||
if True:
|
||
# 1. 首先生成选题
|
||
ai_agent, system_prompt, user_prompt, output_dir = prepare_topic_generation(
|
||
config_file
|
||
)
|
||
|
||
run_id, tweet_topic_record = generate_topics(
|
||
ai_agent, system_prompt, user_prompt, config_file["output_dir"],
|
||
config_file["topic_temperature"], 0.5, 1.5
|
||
)
|
||
|
||
output_dir = os.path.join(config_file["output_dir"], run_id)
|
||
os.makedirs(output_dir, exist_ok=True)
|
||
tweet_topic_record.save_topics(os.path.join(output_dir, "tweet_topic.json"))
|
||
tweet_topic_record.save_prompt(os.path.join(output_dir, "tweet_prompt.txt"))
|
||
# raise Exception("选题生成失败，退出程序")
|
||
if not run_id or not tweet_topic_record:
|
||
print("选题生成失败，退出程序")
|
||
return
|
||
|
||
# 2. 然后生成内容
|
||
print("\n开始根据选题生成内容...")
|
||
|
||
# 加载内容生成的系统提示词
|
||
content_system_prompt = ResourceLoader.load_system_prompt(config_file["content_system_prompt"])
|
||
|
||
|
||
if not content_system_prompt:
|
||
print("内容生成系统提示词为空，使用选题生成的系统提示词")
|
||
content_system_prompt = system_prompt
|
||
|
||
# 直接使用同一个AI Agent实例
|
||
result = generate_content(
|
||
ai_agent, content_system_prompt, tweet_topic_record.topics_list, output_dir, run_id, config_file["prompts_dir"], config_file["resource_dir"],
|
||
config_file["variants"], config_file["content_temperature"]
|
||
)
|
||
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main() |