import os import time from datetime import datetime import argparse import sys import traceback import json from core.ai_agent import AI_Agent # from core.topic_parser import TopicParser # No longer needed directly in main? import core.contentGen as contentGen import core.posterGen as posterGen 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 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", "prompts_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) # 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, tweet_topic_record): """Generates content and posters by calling decoupled functions for each topic.""" if not run_id or not tweet_topic_record or not tweet_topic_record.topics_list: print("Missing run_id or valid topics data. Skipping content and poster generation.") return print("\nStep 2: Processing Topics for Content and Posters...") base_output_dir = config["output_dir"] # --- Initialize shared resources once --- prompt_manager = None ai_agent = None # We might not need to initialize these here if the moved functions handle it # content_gen = None # poster_gen_instance = None initialized_ok = False try: prompt_manager = PromptManager(config) print(f"Initializing AI Agent ({config['model']}) for content/posters...") # --- Read timeout/retry from config --- request_timeout = config.get("request_timeout", 30) # Default 30 seconds max_retries = config.get("max_retries", 3) # Default 3 retries # --- Pass values to AI_Agent --- ai_agent = AI_Agent( config["api_url"], config["model"], config["api_key"], timeout=request_timeout, max_retries=max_retries ) # content_gen = core_contentGen.ContentGenerator() # Instantiation moved to generate_posters_for_topic # poster_gen_instance = core_posterGen.PosterGenerator() # Instantiation moved to generate_posters_for_topic # Check image base directory validity once image_base_dir = config.get("image_base_dir", None) if not image_base_dir or not os.path.isdir(image_base_dir): raise ValueError(f"'image_base_dir' ({image_base_dir}) not specified or not a valid directory in config.") initialized_ok = True except Exception as e: print(f"Error initializing resources for content/poster generation: {e}") traceback.print_exc() if ai_agent: ai_agent.close() return # --- Process each topic using decoupled functions --- total_topics = len(tweet_topic_record.topics_list) for i, topic in enumerate(tweet_topic_record.topics_list): topic_index = i + 1 print(f"\nProcessing Topic {topic_index}/{total_topics}: {topic.get('title', 'N/A')}") # 1. Generate Content for this topic (Call function from utils.tweet_generator) tweet_content_list = generate_content_for_topic( ai_agent, prompt_manager, config, topic, base_output_dir, run_id, topic_index ) # 2. Generate Posters for this topic (Call function from utils.tweet_generator) if tweet_content_list: # Pass only necessary parts of config? Or the whole config? # Pass config for now, the function extracts what it needs. generate_posters_for_topic( config, # Pass config topic, tweet_content_list, base_output_dir, run_id, topic_index # Instances are now created inside generate_posters_for_topic # poster_gen_instance, content_gen, ) else: print(f" Skipping poster generation for Topic {topic_index} due to content generation failure.") # --- Cleanup --- if ai_agent: print("\nClosing shared AI Agent...") ai_agent.close() print("\nFinished processing all topics.") def main(): # No argparse for now, directly load default config config = load_config() # Load from poster_gen_config.json # --- Generate run_id here --- # You can customize this logic later current_run_id = datetime.now().strftime("%Y-%m-%d_%H-%M-%S") print(f"Generated Run ID: {current_run_id}") # --- End run_id generation --- # Execute steps sequentially # Step 1: Generate Topics (pass the generated run_id) run_id_from_step1, tweet_topic_record = run_topic_generation_pipeline(config, run_id=current_run_id) # Step 2: Generate Content and Posters (use the run_id from step 1) # Ensure the returned run_id is used consistently if run_id_from_step1 and tweet_topic_record: generate_content_and_posters_step(config, run_id_from_step1, tweet_topic_record) else: print("Exiting due to issues in topic generation.") if __name__ == "__main__": main()