import os import time from datetime import datetime import argparse import sys import traceback import json import logging from core.ai_agent import AI_Agent # from core.topic_parser import TopicParser # No longer needed directly in main? import utils.content_generator as contentGen import core.poster_gen as poster_gen import core.simple_collage as simple_collage from utils.resource_loader import ResourceLoader from utils.tweet_generator import ( # Import the moved functions run_topic_generation_pipeline, generate_content_for_topic, generate_posters_for_topic ) from utils.prompt_manager import PromptManager # Import PromptManager import random # Import Output Handlers from utils.output_handler import FileSystemOutputHandler, OutputHandler from core.topic_parser import TopicParser from utils.tweet_generator import tweetTopicRecord # Needed only if loading old topics files? # 导入额外配图选择函数 from utils.poster_notes_creator import select_additional_images def load_config(config_path="poster_gen_config.json"): """Loads configuration from a JSON file.""" if not os.path.exists(config_path): print(f"Error: Configuration file '{config_path}' not found.") print("Please copy 'example_config.json' to 'poster_gen_config.json' and customize it.") sys.exit(1) try: with open(config_path, 'r', encoding='utf-8') as f: config = json.load(f) # Basic validation (can be expanded) required_keys = ["api_url", "model", "api_key", "resource_dir", "output_dir", "num", "variants", "topic_system_prompt", "topic_user_prompt", "content_system_prompt", "image_base_dir"] 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"Error: Config file '{config_path}' is missing required keys: {missing_keys}") sys.exit(1) # 验证prompts_dir或prompts_config至少有一个存在 if not ("prompts_dir" in config or "prompts_config" in config): print(f"Error: Config file '{config_path}' must contain either 'prompts_dir' or 'prompts_config'") sys.exit(1) # Resolve relative paths based on config location or a defined base path if necessary # For simplicity, assuming paths in config are relative to project root or absolute return config except json.JSONDecodeError: print(f"Error: Could not decode JSON from '{config_path}'. Check the file format.") sys.exit(1) except Exception as e: print(f"Error loading configuration from '{config_path}': {e}") sys.exit(1) # Removed generate_topics_step function definition from here # Its logic is now in utils.tweet_generator.run_topic_generation_pipeline # --- Main Orchestration Step (Remains in main.py) --- def generate_content_and_posters_step(config, run_id, topics_list, output_handler): """ Step 2: Generates content and posters for each topic in the record. Returns True if successful (at least partially), False otherwise. """ if not topics_list or not topics_list: # print("Skipping content/poster generation: No valid topics found in the record.") logging.warning("Skipping content/poster generation: No valid topics found in the record.") return False # print(f"\n--- Starting Step 2: Content and Poster Generation for run_id: {run_id} ---") logging.info(f"Starting Step 2: Content and Poster Generation for run_id: {run_id}") # print(f"Processing {len(topics_list)} topics...") logging.info(f"Processing {len(topics_list)} topics...") success_flag = False # --- 创建 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"), # 新的配置方式 resource_dir_config=config.get("resource_dir", []), topic_gen_num=config.get("num", 1), # Topic gen num/date used by topic prompts topic_gen_date=config.get("date", ""), content_judger_system_prompt_path=config.get("content_judger_system_prompt") # 添加内容审核系统提示词路径 ) logging.info("PromptManager instance created for Step 2.") except KeyError as e: logging.error(f"Configuration error creating PromptManager: Missing key '{e}'. Cannot proceed with content generation.") return False # --- 结束创建 PromptManager --- ai_agent = None try: # --- Initialize AI Agent for Content Generation --- request_timeout = config.get("request_timeout", 180) # Default 180 seconds max_retries = config.get("max_retries", 3) # Default 3 retries stream_chunk_timeout = config.get("stream_chunk_timeout", 30) # Default 30 seconds for stream chunk ai_agent = AI_Agent( config.get("api_url"), config.get("model"), config.get("api_key"), timeout=request_timeout, max_retries=max_retries, stream_chunk_timeout=stream_chunk_timeout, ) logging.info("AI Agent for content generation initialized.") # --- Iterate through Topics --- for i, topic_item in enumerate(topics_list): topic_index = topic_item.get('index', i + 1) # Use parsed index if available logging.info(f"--- Processing Topic {topic_index}/{len(topics_list)}: {topic_item.get('object', 'N/A')} ---") # Make it stand out # --- Generate Content Variants --- # 读取内容生成需要的参数 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 content_success = generate_content_for_topic( ai_agent, prompt_manager, # Pass PromptManager instance topic_item, run_id, topic_index, output_handler, # Pass OutputHandler instance # 传递具体参数 variants=content_variants, temperature=content_temp, top_p=content_top_p, presence_penalty=content_presence_penalty, enable_content_judge=config.get("enable_content_judge", False) ) # if tweet_content_list: # generate_content_for_topic 现在返回 bool if content_success: logging.info(f"Content generation successful for Topic {topic_index}.") # --- Generate Posters --- # TODO: 重构 generate_posters_for_topic 以移除 config 依赖 # TODO: 需要确定如何将 content 数据传递给 poster 生成步骤 (已解决：函数内部读取) # 临时方案：可能需要在这里读取由 output_handler 保存的 content 文件 # 或者修改 generate_content_for_topic 以返回收集到的 content 数据列表 (选项1) # 暂时跳过 poster 生成的调用，直到确定方案 # --- 重新启用 poster 生成调用 --- logging.info(f"Proceeding to poster generation for Topic {topic_index}...") # --- 读取 Poster 生成所需参数 --- poster_variants = config.get("variants", 1) # 通常与 content variants 相同 poster_assets_dir = config.get("poster_assets_base_dir") img_base_dir = config.get("image_base_dir") res_dir_config = config.get("resource_dir", []) poster_size = tuple(config.get("poster_target_size", [900, 1200])) txt_possibility = config.get("text_possibility", 0.3) img_frame_possibility = config.get("img_frame_possibility", 0.7) text_bg_possibility = config.get("text_bg_possibility", 0) collage_subdir = config.get("output_collage_subdir", "collage_img") poster_subdir = config.get("output_poster_subdir", "poster") poster_filename = config.get("output_poster_filename", "poster.jpg") poster_content_system_prompt = config.get("poster_content_system_prompt", None) collage_style = config.get("collage_style", None) title_possibility = config.get("title_possibility", 0.3) # 检查关键路径是否存在 if not poster_assets_dir or not img_base_dir: logging.error(f"Missing critical paths for poster generation (poster_assets_base_dir or image_base_dir) in config. Skipping posters for topic {topic_index}.") continue # 跳过此主题的海报生成 # --- 结束读取参数 --- with open(poster_content_system_prompt, "r", encoding="utf-8") as f: poster_content_system_prompt = f.read() posters_attempted = generate_posters_for_topic( topic_item=topic_item, output_dir=config["output_dir"], # Base output dir is still needed run_id=run_id, topic_index=topic_index, output_handler=output_handler, # <--- 传递 Output Handler # 传递具体参数 model_name=config["model"], base_url=config["api_url"], api_key=config["api_key"], variants=poster_variants, title_possibility=title_possibility, poster_assets_base_dir=poster_assets_dir, image_base_dir=img_base_dir, resource_dir_config=res_dir_config, poster_target_size=poster_size, text_possibility=txt_possibility, img_frame_possibility=img_frame_possibility, text_bg_possibility=text_bg_possibility, output_collage_subdir=collage_subdir, output_poster_subdir=poster_subdir, output_poster_filename=poster_filename, system_prompt=poster_content_system_prompt, collage_style=collage_style, timeout=request_timeout ) if posters_attempted: logging.info(f"Poster generation process completed for Topic {topic_index}.") # --- 为每个变体添加额外配图 --- logging.info(f"开始为主题 {topic_index} 添加额外配图...") additional_images_count = config.get("additional_images_count", 3) # 循环处理每个变体 for j in range(poster_variants): variant_index = j + 1 variant_dir = os.path.join(config["output_dir"], run_id, f"{topic_index}_{variant_index}") # 获取海报元数据路径 poster_dir = os.path.join(variant_dir, poster_subdir) # 添加目录存在检查 if not os.path.exists(poster_dir): logging.warning(f"海报目录不存在: {poster_dir}, 跳过此变体的额外配图处理") continue try: metadata_files = [f for f in os.listdir(poster_dir) if f.endswith("_metadata.json") and os.path.isfile(os.path.join(poster_dir, f))] except Exception as e: logging.warning(f"访问海报目录时出错: {e}, 跳过此变体的额外配图处理") continue if metadata_files: poster_metadata_path = os.path.join(poster_dir, metadata_files[0]) logging.info(f"为变体 {variant_index} 选择额外配图，使用元数据: {poster_metadata_path}") # 调用额外配图选择函数 try: object_name = topic_item.get("object", "").split(".")[0].replace("景点信息-", "").strip() source_image_dir = os.path.join(img_base_dir, object_name) if os.path.exists(source_image_dir) and os.path.isdir(source_image_dir): # 获取图像选择配置参数 image_selection_config = config.get("image_selection", {}) variation_strength = image_selection_config.get("variation_strength", "medium") extra_effects = image_selection_config.get("extra_effects", True) logging.info(f"图像处理配置: 变化强度={variation_strength}, 额外效果={'启用' if extra_effects else '禁用'}") additional_paths = select_additional_images( run_id=run_id, topic_index=topic_index, variant_index=variant_index, poster_metadata_path=poster_metadata_path, source_image_dir=source_image_dir, num_additional_images=additional_images_count, output_handler=output_handler, variation_strength=variation_strength, extra_effects=extra_effects ) if additional_paths: logging.info(f"已为变体 {variant_index} 选择 {len(additional_paths)} 张额外配图") else: logging.warning(f"未能为变体 {variant_index} 选择任何额外配图") else: logging.warning(f"源图像目录不存在: {source_image_dir}") except Exception as e: logging.exception(f"选择额外配图时发生错误: {e}") else: logging.warning(f"未找到变体 {variant_index} 的海报元数据文件，跳过额外配图选择") # --- 结束添加额外配图 --- success_flag = True # Mark overall success if content AND poster attempts were made else: logging.warning(f"Poster generation skipped or failed early for Topic {topic_index}.") # 即使海报失败，只要内容成功了，也算部分成功？根据需求决定 success_flag # success_flag = True # 取决于是否认为内容成功就足够 # logging.warning(f"Skipping poster generation for Topic {topic_index} pending refactor and data passing strategy.") # Mark overall success if content generation succeeded # success_flag = True else: logging.warning(f"Content generation failed or yielded no valid results for Topic {topic_index}. Skipping posters.") logging.info(f"--- Finished Topic {topic_index} ---") except KeyError as e: # print(f"\nError: Configuration error during content/poster generation: Missing key '{e}'") logging.error(f"Configuration error during content/poster generation: Missing key '{e}'") traceback.print_exc() return False # Indicate failure due to config error except Exception as e: # print(f"\nAn unexpected error occurred during content and poster generation: {e}") # traceback.print_exc() logging.exception("An unexpected error occurred during content and poster generation:") return False # Indicate general failure finally: # Ensure the AI agent is closed if ai_agent: # print("Closing content generation AI Agent...") logging.info("Closing content generation AI Agent...") ai_agent.close() if success_flag: # print("\n--- Step 2: Content and Poster Generation completed (at least partially). ---") logging.info("Step 2: Content and Poster Generation completed (at least partially).") else: # print("\n--- Step 2: Content and Poster Generation completed, but no content/posters were successfully generated or attempted. ---") logging.warning("Step 2: Content and Poster Generation completed, but no content/posters were successfully generated or attempted.") return success_flag # Return True if at least some steps were attempted def main(): # --- Basic Logging Setup --- logging.basicConfig( level=logging.INFO, # Default level format='%(asctime)s - %(levelname)s - [%(filename)s:%(lineno)d] - %(message)s', datefmt='%Y-%m-%d %H:%M:%S' ) # --- End Logging Setup --- parser = argparse.ArgumentParser(description="Travel Content Creator Pipeline") parser.add_argument( "--config", type=str, default="poster_gen_config.json", help="Path to the configuration file (e.g., poster_gen_config.json)" ) parser.add_argument( "--run_id", type=str, default=None, # Default to None, let the pipeline generate it help="Optional specific run ID (e.g., 'my_test_run_01'). If not provided, a timestamp-based ID will be generated." ) parser.add_argument( "--topics_file", type=str, default=None, help="Optional path to a pre-generated topics JSON file. If provided, skips topic generation (Step 1)." ) parser.add_argument( "--debug", action='store_true', # Add debug flag help="Enable debug level logging." ) args = parser.parse_args() # --- Adjust Logging Level if Debug --- if args.debug: logging.getLogger().setLevel(logging.DEBUG) logging.info("Debug logging enabled.") # --- End Debug Level Adjustment --- logging.info("Starting Travel Content Creator Pipeline...") logging.info(f"Using configuration file: {args.config}") if args.run_id: logging.info(f"Using specific run_id: {args.run_id}") if args.topics_file: logging.info(f"Using existing topics file: {args.topics_file}") config = load_config(args.config) if config is None: logging.critical("Failed to load configuration. Exiting.") sys.exit(1) # --- Initialize Output Handler --- # For now, always use FileSystemOutputHandler. Later, this could be configurable. output_handler: OutputHandler = FileSystemOutputHandler(config.get("output_dir", "result")) logging.info(f"Using Output Handler: {output_handler.__class__.__name__}") # --- End Output Handler Init --- run_id = args.run_id # tweet_topic_record = None # No longer the primary way to pass data topics_list = None system_prompt = None user_prompt = None pipeline_start_time = time.time() # --- Step 1: Topic Generation (or Load Existing) --- if args.topics_file: logging.info(f"Skipping Topic Generation (Step 1) - Loading topics from: {args.topics_file}") topics_list = TopicParser.load_topics_from_json(args.topics_file) if topics_list: # Try to infer run_id from filename if not provided if not run_id: try: base = os.path.basename(args.topics_file) # Assuming format "tweet_topic_{run_id}.json" or "tweet_topic.json" if base.startswith("tweet_topic_") and base.endswith(".json"): run_id = base[len("tweet_topic_"): -len(".json")] # print(f"Inferred run_id from topics filename: {run_id}") logging.info(f"Inferred run_id from topics filename: {run_id}") elif base == "tweet_topic.json": # Handle the default name? # Decide how to handle this - maybe use parent dir name or generate new? # For now, let's generate a new one if run_id is still None below logging.warning(f"Loaded topics from default filename '{base}'. Run ID not inferred.") else: # print(f"Warning: Could not infer run_id from topics filename: {base}") logging.warning(f"Could not infer run_id from topics filename: {base}") except Exception as e: # print(f"Warning: Error trying to infer run_id from topics filename: {e}") logging.warning(f"Error trying to infer run_id from topics filename: {e}") # If run_id is still None after trying to infer, generate one if run_id is None: run_id = datetime.now().strftime("%Y-%m-%d_%H-%M-%S_loaded") # print(f"Generated run_id for loaded topics: {run_id}") logging.info(f"Generated run_id for loaded topics: {run_id}") # Prompts are missing when loading from file, handle this if needed later system_prompt = "" # Placeholder user_prompt = "" # Placeholder logging.info(f"Successfully loaded {len(topics_list)} topics for run_id: {run_id}. Prompts are not available.") # Optionally, save the loaded topics using the handler? # output_handler.handle_topic_results(run_id, topics_list, system_prompt, user_prompt) else: # print(f"Error: Failed to load topics from {args.topics_file}. Cannot proceed.") logging.error(f"Failed to load topics from {args.topics_file}. Cannot proceed.") sys.exit(1) else: # print("--- Executing Topic Generation (Step 1) ---") logging.info("Executing Topic Generation (Step 1)...") step1_start = time.time() # Call the updated function, receive raw data 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: # Check if step succeeded # print(f"Step 1 completed successfully in {step1_end - step1_start:.2f} seconds. Run ID: {run_id}") logging.info(f"Step 1 completed successfully in {step1_end - step1_start:.2f} seconds. Run ID: {run_id}") # --- Use Output Handler to save results --- output_handler.handle_topic_results(run_id, topics_list, system_prompt, user_prompt) else: # print("Critical: Topic Generation (Step 1) failed. Exiting.") logging.critical("Topic Generation (Step 1) failed. Exiting.") sys.exit(1) # --- Step 2: Content & Poster Generation --- if run_id is not None and topics_list is not None: # print("\n--- Executing Content and Poster Generation (Step 2) ---") logging.info("Executing Content and Poster Generation (Step 2)...") step2_start = time.time() # TODO: Refactor generate_content_and_posters_step to accept topics_list # and use the output_handler instead of saving files directly. # For now, we might need to pass topics_list and handler, or adapt it. # Let's tentatively adapt the call signature, assuming the function will be refactored. success = generate_content_and_posters_step(config, run_id, topics_list, output_handler) step2_end = time.time() if success: # print(f"Step 2 completed in {step2_end - step2_start:.2f} seconds.") logging.info(f"Step 2 completed in {step2_end - step2_start:.2f} seconds.") else: # print("Warning: Step 2 finished, but may have encountered errors or generated no output.") logging.warning("Step 2 finished, but may have encountered errors or generated no output.") else: # print("Error: Cannot proceed to Step 2: Invalid run_id or topics_list from Step 1.") logging.error("Cannot proceed to Step 2: Invalid run_id or topics_list from Step 1.") # --- Finalize Output --- if run_id: output_handler.finalize(run_id) # --- End Finalize --- pipeline_end_time = time.time() # print(f"\nPipeline finished. Total execution time: {pipeline_end_time - pipeline_start_time:.2f} seconds.") logging.info(f"Pipeline finished. Total execution time: {pipeline_end_time - pipeline_start_time:.2f} seconds.") # print(f"Results for run_id '{run_id}' are in: {os.path.join(config.get('output_dir', 'result'), run_id)}") logging.info(f"Results for run_id '{run_id}' are in: {os.path.join(config.get('output_dir', 'result'), run_id)}") if __name__ == "__main__": main()