TravelContentCreator/utils/resource_loader.py

import os
import random

class ResourceLoader:
    """资源加载器，用于加载提示词和参考资料"""
    
    @staticmethod
    def load_file_content(file_path):
        """加载文件内容"""
        try:
            if os.path.exists(file_path):
                with open(file_path, 'r', encoding='utf-8') as f:
                    content = f.read()
                return content
            else:
                print(f"文件不存在: {file_path}")
                # Return None for non-existent file to distinguish from empty file
                return None 
        except Exception as e:
            print(f"加载文件 '{file_path}' 内容失败: {e}")
            # Return None on error as well
            return None
    
    @staticmethod
    def load_all_refer_files(refer_dir, refer_content_length=50):
        """加载Refer目录下的所有文件内容"""
        refer_content = ""
        if not refer_dir or not os.path.isdir(refer_dir):
             print(f"Warning: Refer directory '{refer_dir}' not found or invalid.")
             return ""
        try:
            files = os.listdir(refer_dir)
            for file in files:
                file_path = os.path.join(refer_dir, file)
                if os.path.isfile(file_path):
                    # Use the updated load_file_content
                    content = ResourceLoader.load_file_content(file_path)
                    if content: # Check if content was loaded successfully
                        # 用\n分割content，取前length条
                        content_lines = content.split("\n")
                        # Ensure refer_content_length doesn't exceed available lines
                        sample_size = min(refer_content_length, len(content_lines))
                        content_lines = random.sample(content_lines, sample_size)
                        content = "\n".join(content_lines)
                        refer_content += f"## {file}\n{content}\n\n"
            return refer_content
        except Exception as e:
            print(f"加载Refer目录文件失败: {e}")
            return ""
    
    @staticmethod
    def find_file_by_name(directory, file_name, exact_match=True):
        """查找文件，支持精确匹配和模糊匹配"""
        if not directory or not file_name:
             return None
        try:
            # 确保传入的文件名包含后缀
            if not file_name.endswith(".txt"):
                file_name = f"{file_name}.txt"
                
            # 精确匹配
            exact_path = os.path.join(directory, file_name)
            if os.path.exists(exact_path) and os.path.isfile(exact_path):
                return exact_path
                
            # 如果不需要精确匹配，尝试模糊匹配
            if not exact_match and os.path.isdir(directory):
                file_name_base = file_name.replace(".txt", "")
                for file in os.listdir(directory):
                     if os.path.isfile(os.path.join(directory, file)) and file_name_base in file:
                         return os.path.join(directory, file)
            
            return None
        except Exception as e:
            print(f"查找文件 '{file_name}' 在 '{directory}' 失败: {e}")
            return None
    
    @staticmethod
    def create_summary_file(output_dir, run_id, topics_count):
        """创建汇总文件并返回路径"""
        summary_file = os.path.join(output_dir, f"summary.md")
        os.makedirs(output_dir, exist_ok=True)
        
        with open(summary_file, 'w', encoding='utf-8') as f:
            f.write(f"# 小红书选题文案生成 - {run_id}\n\n")
            f.write(f"共生成 {topics_count} 篇选题文案\n\n")
            
        return summary_file
        
    @staticmethod
    def update_summary(summary_file, article_index, prompt, result):
        """更新汇总文件"""
        try:
            with open(summary_file, 'a', encoding='utf-8') as f:
                f.write(f"## 文章 {article_index}\n\n")
                f.write(f"### 选题信息\n")
                f.write(f"```\n{prompt}\n```\n\n")
                f.write(f"### 生成内容\n")
                f.write(f"```\n{result}\n```\n\n")
                f.write("--------------------------------\n\n")
        except Exception as e:
            print(f"更新汇总文件时出错: {e}")
         
    @staticmethod
    def save_article(result, prompt, output_dir, run_id, article_index, variant_index):
        """保存生成的文章到文件"""
        try:
            # 确保输出目录存在
            os.makedirs(output_dir, exist_ok=True)
            
            # 创建文件名
            filename = f"article_{article_index}_{variant_index}.txt"
            filepath = os.path.join(output_dir, filename)
            
            # 保存文件
            with open(filepath, 'w', encoding='utf-8') as f:
                f.write(f"prompt: {prompt}\n\n")
                f.write(f"result: {result}\n")
            
            return filepath
        except Exception as e:
            print(f"保存文章时出错: {e}")
            return None