knowledge-agent/tools/structure.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
内容结构化处理模块
简化版本：只提供核心的内容结构化功能
"""

import os
import sys
import json
from typing import Any, Dict, Union

# 导入自定义模块
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from gemini import GeminiProcessor


class StructureProcessor:
    def __init__(self, prompt_file: str = None):
        """
        初始化结构化处理器
        
        Args:
            prompt_file: 提示词文件路径，默认为prompt/structure.md
        """
        self.processor = GeminiProcessor()
        
        # 修复提示词文件路径
        if prompt_file is None:
            # 从tools目录找到项目根目录下的prompt/structure.md
            current_dir = os.path.dirname(os.path.abspath(__file__))  # tools
            project_root = os.path.dirname(current_dir)  # 项目根目录
            prompt_file = os.path.join(project_root, "prompt", "structure.md")
        
        self.prompt_file = prompt_file
        self.system_prompt = self._load_structure_prompt()
    
    def _load_structure_prompt(self) -> str:
        """加载结构化处理提示词"""
        try:
            with open(self.prompt_file, 'r', encoding='utf-8') as f:
                return f.read()
        except Exception as e:
            print(f"加载提示词文件失败: {e}")
            print(f"尝试加载的文件路径: {self.prompt_file}")
            return ""
    
    def process_content(self, content: Union[str, Dict, Any], custom_prompt: str = None) -> str:
        """
        处理内容结构化
        
        Args:
            content: 需要结构化的内容，可以是字符串、字典或JSON对象
            custom_prompt: 自定义提示词，如果为None则使用默认的结构化提示词
            
        Returns:
            str: 结构化后的内容
        """
        try:
            # 如果content是字典，转换为JSON字符串
            if isinstance(content, dict):
                content_str = json.dumps(content, ensure_ascii=False, indent=2)
            else:
                content_str = str(content)
            
            # 使用自定义提示词或默认的结构化提示词
            prompt = custom_prompt or self.system_prompt
            
            # 构建完整的提示词
            full_prompt = f"{prompt}\n\n## 输入\n用户将提供一个包含 `title`、`body_text` 和 `images_comprehension` 的JSON对象。\n\n请处理以下内容：\n{content_str}"
            
            result = self.processor.process(content_str, full_prompt)
            return result
        except Exception as e:
            print(f"内容结构化处理失败: {e}")
            return ""


def main():
    """测试函数"""
    processor = StructureProcessor()
    
    # 测试数据 - 模拟实际的JSON输入
    test_content = {
        "title": "如何制作美味的蛋糕",
        "body_text": "制作蛋糕需要准备一些基本的材料和工具，按照正确的步骤操作就能做出美味的蛋糕。",
        "images_comprehension": [
            "第一步：准备材料 - 面粉、鸡蛋、糖、牛奶",
            "第二步：混合材料 - 将面粉和糖混合，加入鸡蛋和牛奶",
            "第三步：烘烤 - 在180度烤箱中烘烤30分钟",
            "注意事项：确保材料新鲜，烤箱预热到位"
        ]
    }
    
    print("使用默认结构化提示词处理内容...")
    result = processor.process_content(test_content)
    print(f"处理结果:\n{result}")


if __name__ == "__main__":
    main()