TravelContentCreator/utils/tweet_generator.py

#!/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
from TravelContentCreator.utils.resource_loader import ResourceLoader

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, item, prompts_dir, resource_dir, output_dir, run_id, 
                           article_index, variant_index, temperature=0.3, top_p=0.4, presence_penalty=1.5):
    """生成单篇文章内容"""
    try:
        # 构建提示词
        user_prompt = ResourceLoader.build_user_prompt(item, prompts_dir, resource_dir)
        print(f"完成提示词构建，长度为 {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, prompts_dir, resource_dir, 
                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(
    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"
):
    """准备选题生成的环境和参数"""
    # 创建AI Agent
    ai_agent = AI_Agent(base_url, model_name, api_key)
     
    # 加载系统提示词
    with open(system_prompt_path, "r", encoding="utf-8") as f:
        system_prompt = f.read()
    
    # 构建用户提示词
    user_prompt = "你拥有的创作资料如下：\n"
    
    # 加载genPrompts目录下的文件
    gen_prompts_list = os.listdir(gen_prompts_path)
    for gen_prompt_folder in gen_prompts_list:
        folder_path = os.path.join(gen_prompts_path, gen_prompt_folder)
        # 检查是否为目录
        if os.path.isdir(folder_path):
            gen_prompts = os.listdir(folder_path)
            user_prompt += f"{gen_prompt_folder}\n{gen_prompts}\n"
    
    # 加载source文档内的目标文件
    ## 其实只会有Object和Product两个文件夹
    ## 所以可以简化代码
    for dir in resource_dir:
        source_type = dir["type"]
        # source_num = dir["num"]
        source_file_path = dir["file_path"]
        for file in source_file_path:
            with open(file, "r", encoding="utf-8") as f:
                user_prompt += f"{source_type}信息:\n{file.split('/')[-1]}\n{f.read()}\n\n"

    # 加载日期信息
    dateline_path = os.path.join(os.path.dirname(user_prompt_path), "2025各月节日宣传节点时间表.md")
    if os.path.exists(dateline_path):
        with open(dateline_path, "r", encoding="utf-8") as f:
            dateline = f.read()
        user_prompt += f"\n{dateline}"
    
    # 加载用户提示词模板
    with open(user_prompt_path, "r", encoding="utf-8") as f:
        user_prompt += f.read()
    
    # 添加选题数量和日期
    user_prompt += f"\n选题数量：{select_num}\n选题日期：{select_date}\n"
    
    return ai_agent, system_prompt, user_prompt, 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
    # Note: prepare_topic_generation already initializes an AI_Agent
    try:
        ai_agent, system_prompt, user_prompt, base_output_dir = prepare_topic_generation(
            config.get("date", datetime.now().strftime("%Y-%m-%d")), # Use current date if not specified
            config["num"], config["topic_system_prompt"], config["topic_user_prompt"],
            config["api_url"], config["model"], config["api_key"], config["prompts_dir"],
            config["resource_dir"], config["output_dir"]
        )
    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["date"], config_file["num"], config_file["topic_system_prompt"], config_file["topic_user_prompt"],
            config_file["api_url"], config_file["model"], config_file["api_key"], config_file["prompts_dir"], config_file["resource_dir"], config_file["output_dir"]
        )
        
        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()