155 lines
6.6 KiB
Python
155 lines
6.6 KiB
Python
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()
|
|
|
|
|
|
|