rag_server/applications/async_task/record_pattern.py

import json
import traceback
import uuid
import logging
from typing import Dict, List, Tuple, Optional
from dataclasses import dataclass
from applications.utils.mysql import Patterns

# 配置日志
logger = logging.getLogger(__name__)


@dataclass
class ModeData:
    """模式数据类"""

    id: str
    name: str
    percentage: str
    description: str
    detail: str
    output_id: str
    dimension_name: str


@dataclass
class OutputData:
    """产物数据类"""

    id: str
    type: str
    description: str
    content: str
    constraints: str


class RecordPattern:
    def __init__(self, resource):
        self.pattern_manager = Patterns(resource.mysql_client)
        self.milvus_client = resource.milvus_client

    async def record_mode(self, mode: ModeData) -> bool:
        """记录模式数据"""
        try:
            mode_tuple = (
                mode.id,
                mode.name,
                mode.percentage,
                mode.description,
                json.dumps(mode.detail, ensure_ascii=False),
                mode.output_id,
                mode.dimension_name,
            )
            result = await self.pattern_manager.insert_modes([mode_tuple])
            logger.info(f"成功记录模式: {mode.name}")
            return bool(result)
        except Exception as e:
            logger.error(f"记录模式失败: {mode.name}, 错误: {e}")
            print(traceback.format_exc())
            return False

    async def record_output(self, output: OutputData) -> bool:
        """记录产物数据"""
        try:
            output_tuple = (
                output.id,
                output.type,
                output.description,
                output.content,
                output.constraints,
            )
            result = await self.pattern_manager.insert_outputs([output_tuple])
            logger.info(f"成功记录产物: {output.type}")
            return bool(result)
        except Exception as e:
            logger.error(f"记录产物失败: {output.type}, 错误: {e}")
            return False

    @staticmethod
    def _validate_pattern_data(pattern: Dict) -> bool:
        """验证模式数据格式"""
        required_keys = ["维度模式分析"]
        if not all(key in pattern for key in required_keys):
            logger.error(f"模式数据缺少必要字段: {required_keys}")
            return False

        dims = pattern["维度模式分析"]
        if not isinstance(dims, list):
            logger.error("维度模式分析必须是列表类型")
            return False

        for dim in dims:
            if not all(key in dim for key in ["维度名称", "模式列表"]):
                logger.error("维度数据缺少必要字段")
                return False

            for method in dim["模式列表"]:
                required_method_keys = [
                    "模式命名",
                    "模式占比",
                    "模式说明",
                    "分析详情",
                    "可复用产物",
                ]
                if not all(key in method for key in required_method_keys):
                    logger.error(f"模式数据缺少必要字段: {required_method_keys}")
                    return False

                output = method["可复用产物"]
                required_output_keys = ["产物类型", "产物描述", "产物内容", "变量约束"]
                if not all(key in output for key in required_output_keys):
                    logger.error(f"产物数据缺少必要字段: {required_output_keys}")
                    return False

        return True

    @staticmethod
    def _extract_mode_data(method: Dict, dim_name: str) -> Optional[ModeData]:
        """提取模式数据"""
        try:
            mode_id = f"mode-{uuid.uuid4()}"
            output_id = f"output-{uuid.uuid4()}"

            return ModeData(
                id=mode_id,
                name=method.get("模式命名", ""),
                percentage=method.get("模式占比", ""),
                description=method.get("模式说明", ""),
                detail=method.get("分析详情", ""),
                output_id=output_id,
                dimension_name=dim_name,
            )
        except Exception as e:
            logger.error(f"提取模式数据失败: {e}")
            return None

    @staticmethod
    def _extract_output_data(output: Dict, output_id: str) -> Optional[OutputData]:
        """提取产物数据"""
        try:
            return OutputData(
                id=output_id,
                type=output.get("产物类型", ""),
                description=output.get("产物描述", ""),
                content=json.dumps(output.get("产物内容", {}), ensure_ascii=False),
                constraints=json.dumps(output.get("变量约束", {}), ensure_ascii=False),
            )
        except Exception as e:
            logger.error(f"提取产物数据失败: {e}")
            return None

    async def deal(self, pattern: Dict) -> bool:
        """
        处理模式数据

        Args:
            pattern: 模式数据字典

        Returns:
            bool: 处理是否成功
        """
        # 验证数据格式
        if not self._validate_pattern_data(pattern):
            logger.error("模式数据格式验证失败")
            return False

        dims = pattern["维度模式分析"]
        success_count = 0
        total_count = 0

        for dim in dims:
            dim_name = dim["维度名称"]

            for method in dim["模式列表"]:
                total_count += 1

                # 提取模式数据
                mode_data = self._extract_mode_data(method, dim_name)
                if not mode_data:
                    logger.error(f"提取模式数据失败: {dim_name}")
                    continue

                # 提取产物数据
                output = method["可复用产物"]
                output_data = self._extract_output_data(output, mode_data.output_id)
                print(output_data)
                if not output_data:
                    logger.error(f"提取产物数据失败: {mode_data.name}")
                    continue

                # 记录数据
                mode_success = await self.record_mode(mode_data)
                output_success = await self.record_output(output_data)

                if mode_success and output_success:
                    success_count += 1
                    logger.info(f"成功处理模式: {mode_data.name}")
                else:
                    logger.error(f"处理模式失败: {mode_data.name}")

        success_rate = (success_count / total_count) * 100 if total_count > 0 else 0
        logger.info(
            f"模式处理完成: 成功 {success_count}/{total_count} ({success_rate:.1f}%)"
        )

        return success_count > 0