TravelContentCreator/utils/prompt_manager.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Manages the construction of prompts for different AI generation tasks.
"""

import os
import traceback
from .resource_loader import ResourceLoader # Use relative import within the same package

class PromptManager:
    """Handles the loading and construction of prompts."""

    def __init__(self, config):
        """Initializes the PromptManager with the global configuration."""
        self.config = config
        # Instantiate ResourceLoader once, assuming it's mostly static methods or stateless
        self.resource_loader = ResourceLoader()

    def get_topic_prompts(self):
        """Constructs the system and user prompts for topic generation."""
        print("Constructing prompts for topic generation...")
        try:
            # --- System Prompt --- 
            system_prompt_path = self.config.get("topic_system_prompt")
            if not system_prompt_path:
                print("Error: topic_system_prompt path not specified in config.")
                return None, None
                
            # Use ResourceLoader's static method directly
            system_prompt = ResourceLoader.load_file_content(system_prompt_path)
            if not system_prompt:
                 print(f"Error: Failed to load topic system prompt from '{system_prompt_path}'.")
                 return None, None

            # --- User Prompt --- 
            user_prompt_path = self.config.get("topic_user_prompt")
            if not user_prompt_path:
                print("Error: topic_user_prompt path not specified in config.")
                return None, None

            base_user_prompt = ResourceLoader.load_file_content(user_prompt_path)
            if base_user_prompt is None: # Check for None explicitly
                 print(f"Error: Failed to load base topic user prompt from '{user_prompt_path}'.")
                 return None, None
            
            # --- Build the dynamic part of the user prompt (Logic moved from prepare_topic_generation) ---
            user_prompt_dynamic = "你拥有的创作资料如下：\n"
            
            # Add genPrompts directory structure
            gen_prompts_path = self.config.get("prompts_dir")
            if gen_prompts_path and os.path.isdir(gen_prompts_path):
                try:
                    gen_prompts_list = os.listdir(gen_prompts_path)
                    for gen_prompt_folder in gen_prompts_list:
                        folder_path = os.path.join(gen_prompts_path, gen_prompt_folder)
                        if os.path.isdir(folder_path):
                            try:
                                # List files, filter out subdirs if needed
                                gen_prompts_files = [f for f in os.listdir(folder_path) if os.path.isfile(os.path.join(folder_path, f))]
                                user_prompt_dynamic += f"{gen_prompt_folder}\n{gen_prompts_files}\n"
                            except OSError as e:
                                print(f"Warning: Could not list directory {folder_path}: {e}")
                except OSError as e:
                    print(f"Warning: Could not list base prompts directory {gen_prompts_path}: {e}")
            else:
                 print(f"Warning: Prompts directory '{gen_prompts_path}' not found or invalid.")

            # Add resource directory contents
            resource_dir_config = self.config.get("resource_dir", [])
            for dir_info in resource_dir_config:
                source_type = dir_info.get("type", "UnknownType")
                source_file_paths = dir_info.get("file_path", [])
                for file_path in source_file_paths:
                    # Use ResourceLoader's static method
                    file_content = ResourceLoader.load_file_content(file_path)
                    if file_content:
                         user_prompt_dynamic += f"{source_type}信息:\n{os.path.basename(file_path)}\n{file_content}\n\n"
                    else:
                         print(f"Warning: Could not load resource file {file_path}")

            # Add dateline information (optional)
            user_prompt_dir = os.path.dirname(user_prompt_path)
            dateline_path = os.path.join(user_prompt_dir, "2025各月节日宣传节点时间表.md") # Consider making this configurable
            if os.path.exists(dateline_path):
                dateline_content = ResourceLoader.load_file_content(dateline_path)
                if dateline_content:
                     user_prompt_dynamic += f"\n{dateline_content}"
            
            # Combine dynamic part, base template, and final parameters
            user_prompt = user_prompt_dynamic + base_user_prompt
            
            select_num = self.config.get("num", 1) # Default to 1 if not specified
            select_date = self.config.get("date", "")
            user_prompt += f"\n选题数量：{select_num}\n选题日期：{select_date}\n"
            # --- End of moved logic --- 
            
            print(f"Topic prompts constructed. System: {len(system_prompt)} chars, User: {len(user_prompt)} chars.")
            return system_prompt, user_prompt

        except Exception as e:
            print(f"Error constructing topic prompts: {e}")
            traceback.print_exc()
            return None, None

    def get_content_prompts(self, topic_item):
        """Constructs the system and user prompts for content generation based on a topic item."""
        print(f"Constructing content prompts for topic: {topic_item.get('object', 'N/A')}...")
        try:
            # --- System Prompt --- 
            system_prompt_path = self.config.get("content_system_prompt")
            if not system_prompt_path:
                 print("Error: content_system_prompt path not specified in config.")
                 return None, None
            # Use ResourceLoader's static method. load_system_prompt was just load_file_content.
            system_prompt = ResourceLoader.load_file_content(system_prompt_path)
            if not system_prompt:
                 print(f"Error: Failed to load content system prompt from '{system_prompt_path}'.")
                 return None, None

            # --- User Prompt (Logic moved from ResourceLoader.build_user_prompt) ---
            user_prompt = ""
            prompts_dir = self.config.get("prompts_dir")
            resource_dir_config = self.config.get("resource_dir", [])

            if not prompts_dir or not os.path.isdir(prompts_dir):
                print(f"Warning: Prompts directory '{prompts_dir}' not found or invalid. Content user prompt might be incomplete.")
                # Decide whether to return error or continue with potentially incomplete prompt

            # 1. 添加Demand部分 (直接使用 topic_item['logic'] 的描述性文本)
            try:
                demand_description = topic_item.get('logic')
                if demand_description:
                    user_prompt += f"Demand Logic:\n{demand_description}\n"
                else:
                    print("Warning: 'logic' key missing or empty in topic_item for Demand prompt.")
            except Exception as e:
                 print(f"Error processing Demand description: {e}")

            # 2. 添加Object信息 (based on topic_item['object'])
            try:
                object_name_base = topic_item['object'] # This might be '景点信息-XXX.txt'
                object_file_path = None
                # Find the full path for the object file from config
                for dir_info in resource_dir_config:
                    if dir_info.get("type") == "Object":
                        for file_path in dir_info.get("file_path", []):
                            # Match basename, assuming topic_item['object'] is the basename
                            # if os.path.basename(file_path) == object_name_base:
                            # Use containment check instead of exact match
                            if object_name_base in os.path.basename(file_path):
                                 object_file_path = file_path
                                 break
                    if object_file_path: break
                
                if object_file_path:
                     object_content = ResourceLoader.load_file_content(object_file_path)
                     if object_content:
                         user_prompt += f"Object信息:\n{object_content}\n"
                     else:
                         print(f"Warning: Object file could not be loaded: {object_file_path}")
                else:
                     # If basename match fails, maybe topic_item['object'] is just 'XXX'?
                     # Try finding based on substring? This might be ambiguous.
                     print(f"Warning: Object file path not found in config matching basename: {object_name_base}")
            except KeyError:
                 print("Warning: 'object' key missing in topic_item for Object prompt.")
            except Exception as e:
                 print(f"Error processing Object prompt: {e}")

            # 3. 添加Product信息 (if applicable)
            try:
                product_name = topic_item.get('product')
                product_logic_description = topic_item.get('product_logic') # Directly use this description
                if product_name:
                    # Add Product Logic description first (if available)
                    if product_logic_description:
                        user_prompt += f"Product Logic:\n{product_logic_description}\n"
                    else:
                        print(f"Warning: 'product_logic' key missing or empty for product '{product_name}'.")

                    # Then, load Product Info file
                    product_file_path = None
                    for dir_info in resource_dir_config:
                        if dir_info.get("type") == "Product":
                            for file_path in dir_info.get("file_path", []):
                                if product_name in os.path.basename(file_path):
                                    product_file_path = file_path
                                    break
                        if product_file_path: break
                    
                    if product_file_path:
                         product_content = ResourceLoader.load_file_content(product_file_path)
                         if product_content:
                             user_prompt += f"Product Info:\n{product_content}\n"
                         else:
                             print(f"Warning: Product file could not be loaded: {product_file_path}")
                    else:
                         print(f"Warning: Product file path not found in config for: {product_name}")
                # Removed KeyError check for product_logic as it's now optional text
            except KeyError as e:
                 print(f"Warning: Missing 'product' key in topic_item for Product prompt: {e}") # Should only warn for product now
            except Exception as e:
                 print(f"Error processing Product prompt: {e}")

            # 4. 添加Style信息 (加载文件 based on topic_item['style'] from Style/ folder)
            try:
                style_filename = topic_item.get('style')
                if style_filename:
                    # Ensure .txt extension
                    if not style_filename.lower().endswith('.txt'):
                        style_file = f"{style_filename}.txt"
                    else:
                        style_file = style_filename
                    style_path = os.path.join(prompts_dir, "Style", style_file) # Look in Style/ subfolder
                    style_content = ResourceLoader.load_file_content(style_path)
                    if style_content:
                        user_prompt += f"Style Info:\n{style_content}\n" # Changed label for clarity
                    else:
                        print(f"Warning: Style file not found or empty: {style_path}")
                else:
                     print("Warning: 'style' key missing or empty in topic_item.")
            except KeyError:
                 print("Warning: 'style' key missing in topic_item for Style prompt.") # Keep KeyError check just in case
            except Exception as e:
                 print(f"Error processing Style prompt: {e}")

            # 5. 添加Target Audience信息 (加载文件 based on topic_item['target_audience'] from Demand/ folder)
            try:
                target_audience_filename = topic_item.get('target_audience')
                if target_audience_filename:
                    # Ensure .txt extension (or use the check as done for style)
                    if not target_audience_filename.lower().endswith('.txt'):
                        target_audience_file = f"{target_audience_filename}.txt"
                    else:
                        target_audience_file = target_audience_filename
                    # Assuming target audience files are in the 'Demand' subdirectory based on old prompt
                    # Load from Demand/ subfolder as confirmed
                    target_audience_path = os.path.join(prompts_dir, "Demand", target_audience_file) 
                    target_audience_content = ResourceLoader.load_file_content(target_audience_path)
                    if target_audience_content:
                        user_prompt += f"Target Audience Info:\n{target_audience_content}\n"
                    else:
                        print(f"Warning: Target Audience file not found or empty: {target_audience_path}")
                else:
                     print("Warning: 'target_audience' key missing or empty in topic_item.")
            except Exception as e:
                 print(f"Error processing Target Audience prompt: {e}")

            # 6. 添加Refer信息 (加载 Refer/ 目录下所有文件的内容)
            try:
                refer_dir = os.path.join(prompts_dir, "Refer")
                if os.path.isdir(refer_dir):
                    refer_content_all = ""
                    refer_files = [f for f in os.listdir(refer_dir) if os.path.isfile(os.path.join(refer_dir, f))]
                    print(f"Found {len(refer_files)} files in Refer directory: {refer_files}")
                    for refer_file in refer_files:
                        refer_path = os.path.join(refer_dir, refer_file)
                        content = ResourceLoader.load_file_content(refer_path)
                        if content:
                            refer_content_all += f"--- Refer File: {refer_file} ---\n{content}\n\n"
                        else:
                            print(f"Warning: Could not load Refer file: {refer_path}")
                    if refer_content_all:
                         user_prompt += f"Refer Info:\n{refer_content_all}"
                    else:
                         print("Warning: No content loaded from Refer directory.")
                else:
                     print(f"Warning: Refer directory not found: {refer_dir}")
            except Exception as e:
                 print(f"Error processing Refer directory: {e}")

            # --- End of prompt construction logic ---
            
            print(f"Content prompts constructed. System: {len(system_prompt)} chars, User: {len(user_prompt)} chars.")
            return system_prompt, user_prompt

        except KeyError as e:
             # Catch potential KeyErrors from accessing topic_item if a required key is missing early on
             print(f"Error constructing content prompts: Missing essential key '{e}' in topic_item: {topic_item}")
             return None, None
        except Exception as e:
            print(f"Error constructing content prompts: {e}")
            traceback.print_exc()
            return None, None