145 lines
6.9 KiB
Python
145 lines
6.9 KiB
Python
import os
|
|
import json
|
|
import logging
|
|
from abc import ABC, abstractmethod
|
|
|
|
class OutputHandler(ABC):
|
|
"""Abstract base class for handling the output of the generation pipeline."""
|
|
|
|
@abstractmethod
|
|
def handle_topic_results(self, run_id: str, topics_list: list, system_prompt: str, user_prompt: str):
|
|
"""Handles the results from the topic generation step."""
|
|
pass
|
|
|
|
@abstractmethod
|
|
def handle_content_variant(self, run_id: str, topic_index: int, variant_index: int, content_data: dict, prompt_data: str):
|
|
"""Handles the results for a single content variant."""
|
|
pass
|
|
|
|
@abstractmethod
|
|
def handle_poster_configs(self, run_id: str, topic_index: int, config_data: list | dict):
|
|
"""Handles the poster configuration generated for a topic."""
|
|
pass
|
|
|
|
@abstractmethod
|
|
def handle_generated_image(self, run_id: str, topic_index: int, variant_index: int, image_type: str, image_data, output_filename: str):
|
|
"""Handles a generated image (collage or final poster).
|
|
|
|
Args:
|
|
image_type: Either 'collage' or 'poster'.
|
|
image_data: The image data (e.g., PIL Image object or bytes).
|
|
output_filename: The desired filename for the output (e.g., 'poster.jpg').
|
|
"""
|
|
pass
|
|
|
|
@abstractmethod
|
|
def finalize(self, run_id: str):
|
|
"""Perform any final actions for the run (e.g., close files, upload manifests)."""
|
|
pass
|
|
|
|
|
|
class FileSystemOutputHandler(OutputHandler):
|
|
"""Handles output by saving results to the local file system."""
|
|
|
|
def __init__(self, base_output_dir: str = "result"):
|
|
self.base_output_dir = base_output_dir
|
|
logging.info(f"FileSystemOutputHandler initialized. Base output directory: {self.base_output_dir}")
|
|
|
|
def _get_run_dir(self, run_id: str) -> str:
|
|
"""Gets the specific directory for a run, creating it if necessary."""
|
|
run_dir = os.path.join(self.base_output_dir, run_id)
|
|
os.makedirs(run_dir, exist_ok=True)
|
|
return run_dir
|
|
|
|
def _get_variant_dir(self, run_id: str, topic_index: int, variant_index: int, subdir: str | None = None) -> str:
|
|
"""Gets the specific directory for a variant, optionally within a subdirectory (e.g., 'poster'), creating it if necessary."""
|
|
run_dir = self._get_run_dir(run_id)
|
|
variant_base_dir = os.path.join(run_dir, f"{topic_index}_{variant_index}")
|
|
target_dir = variant_base_dir
|
|
if subdir:
|
|
target_dir = os.path.join(variant_base_dir, subdir)
|
|
os.makedirs(target_dir, exist_ok=True)
|
|
return target_dir
|
|
|
|
def handle_topic_results(self, run_id: str, topics_list: list, system_prompt: str, user_prompt: str):
|
|
run_dir = self._get_run_dir(run_id)
|
|
|
|
# Save topics list
|
|
topics_path = os.path.join(run_dir, f"tweet_topic_{run_id}.json")
|
|
try:
|
|
with open(topics_path, "w", encoding="utf-8") as f:
|
|
json.dump(topics_list, f, ensure_ascii=False, indent=4)
|
|
logging.info(f"Topics list saved successfully to: {topics_path}")
|
|
except Exception as e:
|
|
logging.exception(f"Error saving topic JSON file to {topics_path}:")
|
|
|
|
# Save prompts
|
|
prompt_path = os.path.join(run_dir, f"topic_prompt_{run_id}.txt")
|
|
try:
|
|
with open(prompt_path, "w", encoding="utf-8") as f:
|
|
f.write("--- SYSTEM PROMPT ---\n")
|
|
f.write(system_prompt + "\n\n")
|
|
f.write("--- USER PROMPT ---\n")
|
|
f.write(user_prompt + "\n")
|
|
logging.info(f"Topic prompts saved successfully to: {prompt_path}")
|
|
except Exception as e:
|
|
logging.exception(f"Error saving topic prompts file to {prompt_path}:")
|
|
|
|
def handle_content_variant(self, run_id: str, topic_index: int, variant_index: int, content_data: dict, prompt_data: str):
|
|
"""Saves content JSON and prompt for a specific variant."""
|
|
variant_dir = self._get_variant_dir(run_id, topic_index, variant_index)
|
|
|
|
# Save content JSON
|
|
content_path = os.path.join(variant_dir, "article.json")
|
|
try:
|
|
with open(content_path, "w", encoding="utf-8") as f:
|
|
json.dump(content_data, f, ensure_ascii=False, indent=4)
|
|
logging.info(f"Content JSON saved to: {content_path}")
|
|
except Exception as e:
|
|
logging.exception(f"Failed to save content JSON to {content_path}: {e}")
|
|
|
|
# Save content prompt
|
|
prompt_path = os.path.join(variant_dir, "tweet_prompt.txt")
|
|
try:
|
|
with open(prompt_path, "w", encoding="utf-8") as f:
|
|
# Assuming prompt_data is the user prompt used for this variant
|
|
f.write(prompt_data + "\n")
|
|
logging.info(f"Content prompt saved to: {prompt_path}")
|
|
except Exception as e:
|
|
logging.exception(f"Failed to save content prompt to {prompt_path}: {e}")
|
|
|
|
def handle_poster_configs(self, run_id: str, topic_index: int, config_data: list | dict):
|
|
"""Saves the complete poster configuration list/dict for a topic."""
|
|
run_dir = self._get_run_dir(run_id)
|
|
config_path = os.path.join(run_dir, f"topic_{topic_index}_poster_configs.json")
|
|
try:
|
|
with open(config_path, 'w', encoding='utf-8') as f_cfg_topic:
|
|
json.dump(config_data, f_cfg_topic, ensure_ascii=False, indent=4)
|
|
logging.info(f"Saved complete poster configurations for topic {topic_index} to: {config_path}")
|
|
except Exception as save_err:
|
|
logging.error(f"Failed to save complete poster configurations for topic {topic_index} to {config_path}: {save_err}")
|
|
|
|
def handle_generated_image(self, run_id: str, topic_index: int, variant_index: int, image_type: str, image_data, output_filename: str):
|
|
"""Saves a generated image (PIL Image) to the appropriate variant subdirectory."""
|
|
subdir = None
|
|
if image_type == 'collage':
|
|
subdir = 'collage_img' # TODO: Make these subdir names configurable?
|
|
elif image_type == 'poster':
|
|
subdir = 'poster'
|
|
else:
|
|
logging.warning(f"Unknown image type '{image_type}'. Saving to variant root.")
|
|
subdir = None # Save directly in variant dir if type is unknown
|
|
|
|
target_dir = self._get_variant_dir(run_id, topic_index, variant_index, subdir=subdir)
|
|
save_path = os.path.join(target_dir, output_filename)
|
|
|
|
try:
|
|
# Assuming image_data is a PIL Image object based on posterGen/simple_collage
|
|
image_data.save(save_path)
|
|
logging.info(f"Saved {image_type} image to: {save_path}")
|
|
except Exception as e:
|
|
logging.exception(f"Failed to save {image_type} image to {save_path}: {e}")
|
|
|
|
def finalize(self, run_id: str):
|
|
logging.info(f"FileSystemOutputHandler finalizing run: {run_id}. No specific actions needed.")
|
|
pass # Nothing specific to do for file system finalize |