TravelContentCreator/utils/tweet/content_judger.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
内容审核模块
"""

import logging
import json
from json_repair import loads as json_repair_loads
from typing import Dict, Any

from core.ai import AIAgent
from core.config import ConfigManager, GenerateTopicConfig
from utils.prompts import JudgerPromptBuilder

logger = logging.getLogger(__name__)


class ContentJudger:
    """内容审核类，使用AI评估和修正内容"""

    def __init__(self, ai_agent: AIAgent, config_manager: ConfigManager):
        """
        初始化内容审核器

        Args:
            ai_agent: AIAgent实例
            config_manager: 配置管理器
        """
        self.ai_agent = ai_agent
        self.config: GenerateTopicConfig = config_manager.get_config('topic_gen', GenerateTopicConfig)
        self.prompt_builder = JudgerPromptBuilder(config_manager)

    async def judge_content(self, generated_content: str, topic: Dict[str, Any]) -> Dict[str, Any]:
        """
        调用AI审核生成的内容

        Args:
            generated_content: 已生成的原始内容（JSON字符串格式）
            topic: 与内容相关的原始选题字典

        Returns:
            一个包含审核结果的字典
        """
        logger.info("开始审核生成的内容...")
        
        # 1. 构建提示
        system_prompt = self.prompt_builder.get_system_prompt()
        user_prompt = self.prompt_builder.build_user_prompt(
            generated_content=generated_content,
            topic=topic
        )

        # 2. 调用AI进行审核
        try:
            raw_result, _, _, _ = await self.ai_agent.generate_text(
                system_prompt=system_prompt,
                user_prompt=user_prompt,
                use_stream=False
            )
        except Exception as e:
            logger.critical(f"内容审核时AI调用失败: {e}", exc_info=True)
            return {"judge_success": False, "error": str(e)}

        # 3. 解析结果
        try:
            judged_data = json_repair_loads(raw_result)
            if isinstance(judged_data, dict) and "title" in judged_data and "content" in judged_data:
                judged_data["judge_success"] = True
                logger.info("内容审核成功完成。")
                return judged_data
            else:
                logger.warning(f"审核响应JSON格式不正确或缺少键: {judged_data}")
                return {"judge_success": False, "error": "Invalid JSON response", "raw_response": raw_result}
        except (json.JSONDecodeError, ValueError) as e:
            logger.error(f"解析审核响应JSON失败: {e}")
            return {"judge_success": False, "error": "JSONDecodeError", "raw_response": raw_result}