238 lines
11 KiB
Python
238 lines
11 KiB
Python
#!/usr/bin/env python3
|
|
# -*- coding: utf-8 -*-
|
|
|
|
import os
|
|
import sys
|
|
import json
|
|
import time
|
|
from datetime import datetime
|
|
import logging
|
|
|
|
# --- Setup Project Root Path ---
|
|
PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
|
if PROJECT_ROOT not in sys.path:
|
|
sys.path.append(PROJECT_ROOT)
|
|
# --- End Setup ---
|
|
|
|
# --- Imports from the project ---
|
|
try:
|
|
from utils.tweet_generator import run_topic_generation_pipeline, generate_content_for_topic, generate_posters_for_topic
|
|
from core.topic_parser import TopicParser
|
|
from utils.prompt_manager import PromptManager # Needed for content generation
|
|
from core.ai_agent import AI_Agent # Needed for content generation
|
|
# from utils.tweet_generator import tweetTopicRecord # No longer needed directly
|
|
from utils.output_handler import FileSystemOutputHandler, OutputHandler # Import handlers
|
|
except ImportError as e:
|
|
logging.critical(f"ImportError: {e}. Ensure all core/utils modules are available and '{PROJECT_ROOT}' is in sys.path.")
|
|
sys.exit(1)
|
|
# --- End Imports ---
|
|
|
|
def load_config(config_path):
|
|
"""Loads configuration from a JSON file."""
|
|
try:
|
|
with open(config_path, 'r', encoding='utf-8') as f:
|
|
config = json.load(f)
|
|
logging.info(f"Config loaded successfully from {config_path}")
|
|
return config
|
|
except FileNotFoundError:
|
|
logging.error(f"Error: Configuration file not found at {config_path}")
|
|
return None
|
|
except json.JSONDecodeError:
|
|
logging.error(f"Error: Could not decode JSON from {config_path}")
|
|
return None
|
|
except Exception as e:
|
|
logging.exception(f"An unexpected error occurred loading config {config_path}:")
|
|
return None
|
|
|
|
def main_test():
|
|
# --- Basic Logging Setup ---
|
|
logging.basicConfig(
|
|
level=logging.INFO,
|
|
format='%(asctime)s - %(levelname)s - [%(filename)s:%(lineno)d] - %(message)s',
|
|
datefmt='%Y-%m-%d %H:%M:%S'
|
|
)
|
|
# --- End Logging Setup ---
|
|
|
|
logging.info("--- Starting Pipeline Steps Test ---")
|
|
|
|
# 1. Load Configuration
|
|
config_path = os.path.join(PROJECT_ROOT, "poster_gen_config.json") # Use project root path
|
|
config = load_config(config_path)
|
|
if config is None:
|
|
logging.critical("Failed to load configuration. Exiting test.")
|
|
sys.exit(1)
|
|
|
|
# --- Initialize Output Handler ---
|
|
output_handler: OutputHandler = FileSystemOutputHandler(config.get("output_dir", "result"))
|
|
logging.info(f"Using Output Handler: {output_handler.__class__.__name__}")
|
|
# --- End Output Handler Init ---
|
|
|
|
# 2. Define a unique Run ID for this test run
|
|
test_run_id = f"test_pipeline_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
|
|
logging.info(f"Using Run ID: {test_run_id}")
|
|
|
|
# --- Step 1: Topic Generation ---
|
|
logging.info("\n--- Testing Step 1: Topic Generation ---")
|
|
step1_start = time.time()
|
|
# run_topic_generation_pipeline still takes config internally
|
|
run_id_step1, topics_list, system_prompt, user_prompt = run_topic_generation_pipeline(config, run_id=test_run_id)
|
|
step1_end = time.time()
|
|
|
|
if run_id_step1 is None or topics_list is None:
|
|
logging.error("Topic generation (Step 1) failed. Cannot proceed.")
|
|
sys.exit(1)
|
|
if run_id_step1 != test_run_id:
|
|
logging.warning(f"run_id returned from step 1 ({run_id_step1}) differs from the one provided ({test_run_id}). Using the returned one.")
|
|
test_run_id = run_id_step1 # Use the ID actually used by the pipeline
|
|
|
|
logging.info(f"Step 1 finished in {step1_end - step1_start:.2f} seconds.")
|
|
logging.info(f"Generated {len(topics_list)} topics for Run ID: {test_run_id}")
|
|
# Use output handler to save topic results (mimics main.py)
|
|
output_handler.handle_topic_results(test_run_id, topics_list, system_prompt, user_prompt)
|
|
|
|
# --- Step 2: Content and Poster Generation ---
|
|
logging.info("\n--- Testing Step 2: Content and Poster Generation ---")
|
|
step2_start = time.time()
|
|
|
|
# Initialize resources needed for step 2, extracting params from config
|
|
prompt_manager = None
|
|
ai_agent = None
|
|
step2_success_flag = False
|
|
try:
|
|
# --- Create PromptManager ---
|
|
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_dir=config.get("prompts_dir"),
|
|
resource_dir_config=config.get("resource_dir", []),
|
|
topic_gen_num=config.get("num", 1),
|
|
topic_gen_date=config.get("date", "")
|
|
)
|
|
logging.info("PromptManager instance created for Step 2 test.")
|
|
|
|
# --- Create AI Agent ---
|
|
ai_api_url = config.get("api_url")
|
|
ai_model = config.get("model")
|
|
ai_api_key = config.get("api_key")
|
|
request_timeout = config.get("request_timeout", 30)
|
|
max_retries = config.get("max_retries", 3)
|
|
stream_chunk_timeout = config.get("stream_chunk_timeout", 60) # Read stream timeout
|
|
if not all([ai_api_url, ai_model, ai_api_key]):
|
|
raise ValueError("Missing required AI configuration (api_url, model, api_key)")
|
|
logging.info("Initializing AI Agent for content generation test...")
|
|
ai_agent = AI_Agent(
|
|
base_url=ai_api_url,
|
|
model_name=ai_model,
|
|
api=ai_api_key,
|
|
timeout=request_timeout,
|
|
max_retries=max_retries,
|
|
stream_chunk_timeout=stream_chunk_timeout # Pass stream timeout
|
|
)
|
|
|
|
total_topics = len(topics_list)
|
|
logging.info(f"Processing {total_topics} topics for content/posters...")
|
|
for i, topic_item in enumerate(topics_list):
|
|
topic_index = topic_item.get('index', i + 1)
|
|
logging.info(f"--- Processing Topic {topic_index}/{total_topics} ---")
|
|
|
|
# --- Generate Content ---
|
|
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)
|
|
|
|
content_success = generate_content_for_topic(
|
|
ai_agent=ai_agent,
|
|
prompt_manager=prompt_manager,
|
|
topic_item=topic_item,
|
|
run_id=test_run_id,
|
|
topic_index=topic_index,
|
|
output_handler=output_handler,
|
|
variants=content_variants,
|
|
temperature=content_temp,
|
|
top_p=content_top_p,
|
|
presence_penalty=content_presence_penalty
|
|
)
|
|
|
|
if content_success:
|
|
logging.info(f"Content generation successful for Topic {topic_index}.")
|
|
# --- Generate Posters ---
|
|
poster_variants = config.get("variants", 1)
|
|
poster_assets_dir = config.get("poster_assets_base_dir")
|
|
img_base_dir = config.get("image_base_dir")
|
|
mod_img_subdir = config.get("modify_image_subdir", "modify")
|
|
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)
|
|
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")
|
|
cam_img_subdir = config.get("camera_image_subdir", "相机")
|
|
|
|
if not poster_assets_dir or not img_base_dir:
|
|
logging.error(f"Missing critical paths for poster generation. Skipping posters for topic {topic_index}.")
|
|
continue
|
|
|
|
posters_attempted = generate_posters_for_topic(
|
|
topic_item=topic_item,
|
|
output_dir=config.get("output_dir", "result"),
|
|
run_id=test_run_id,
|
|
topic_index=topic_index,
|
|
output_handler=output_handler,
|
|
variants=poster_variants,
|
|
poster_assets_base_dir=poster_assets_dir,
|
|
image_base_dir=img_base_dir,
|
|
modify_image_subdir=mod_img_subdir,
|
|
resource_dir_config=res_dir_config,
|
|
poster_target_size=poster_size,
|
|
text_possibility=txt_possibility,
|
|
output_collage_subdir=collage_subdir,
|
|
output_poster_subdir=poster_subdir,
|
|
output_poster_filename=poster_filename,
|
|
camera_image_subdir=cam_img_subdir
|
|
)
|
|
if posters_attempted:
|
|
logging.info(f"Poster generation process completed for Topic {topic_index}.")
|
|
step2_success_flag = True
|
|
else:
|
|
logging.warning(f"Poster generation skipped/failed for Topic {topic_index}.")
|
|
else:
|
|
logging.warning(f"Content generation failed for Topic {topic_index}. Skipping posters.")
|
|
logging.info(f"--- Finished Topic {topic_index} ---")
|
|
|
|
except ValueError as e:
|
|
logging.error(f"Configuration error during Step 2 setup: {e}")
|
|
step2_success_flag = False
|
|
except Exception as e:
|
|
logging.exception("An error occurred during Step 2 processing:")
|
|
step2_success_flag = False # Ensure flag is false on error
|
|
finally:
|
|
if ai_agent:
|
|
logging.info("Closing AI Agent for content generation test...")
|
|
ai_agent.close()
|
|
# --- End Simulated Step 2 Logic ---
|
|
|
|
step2_end = time.time()
|
|
if step2_success_flag:
|
|
logging.info(f"Step 2 finished in {step2_end - step2_start:.2f} seconds.")
|
|
else:
|
|
logging.warning(f"Step 2 finished in {step2_end - step2_start:.2f} seconds, but encountered errors or generated no output.")
|
|
|
|
# --- Finalize Output ---
|
|
if test_run_id:
|
|
output_handler.finalize(test_run_id)
|
|
# --- End Finalize ---
|
|
|
|
# --- Test Summary ---
|
|
logging.info("\n--- Pipeline Steps Test Summary ---")
|
|
logging.info(f"Run ID: {test_run_id}")
|
|
output_location = os.path.join(config["output_dir"], test_run_id)
|
|
logging.info(f"Check output files in: {output_location}")
|
|
if os.path.exists(output_location):
|
|
logging.info("Output directory exists.")
|
|
else:
|
|
logging.warning("Output directory NOT found.")
|
|
|
|
if __name__ == "__main__":
|
|
main_test() |