QueryRewrite/rag2_0/intent_recognition/DataModels.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
File: DataModels.py
Author: oyyz
Date: 2025-05-13
Description: 提取和分类的数据模型
"""

from pydantic import BaseModel, Field
from typing import List, Optional, Dict, Tuple
from enum import Enum

class SoftwareName(str, Enum):
    """软件名称枚举类"""
    D3 = "配网工程计价通D3软件"
    C1 = "新型储能电站建设计价通C1软件"
    Z1 = "西藏电力工程计价通Z1软件"
    T1 = "技改检修工程计价通T1软件"
    T1_LIST = "技改检修清单计价通T1软件"
    MAIN = "主网电力建设计价通软件"
    UNKNOWN = "" # 未知

    # 软件别名映射
    ALIASES = {
        D3: "别名包括：配网D3、D3软件、配网工程软件等其他类似称呼",
        C1: "别名包括：储能C1、C1软件、储能电站软件、储能软件等其他类似称呼",
        Z1: "别名包括：西藏Z1、Z1软件、西藏电力软件等其他类似称呼",
        T1: "别名包括：技改T1、T1软件、技改检修软件等其他类似称呼",
        T1_LIST: "别名包括：技改清单T1、T1清单软件、技改检修清单软件等其他类似称呼",
        MAIN: "别名包括：主网软件、电力建设软件、主网建设软件、博微电力建设计价通等其他类似称呼"
    }

# 定义输出模型
class Term(BaseModel):
    name: str = Field(description="专业名词")
    synonymous: List[str] = Field(description="同义词列表")
    description: str = Field(description="描述信息", default="")
    
    def __hash__(self):
        return hash(self.name)
    
    def __eq__(self, other):
        if isinstance(other, Term):
            return self.name == other.name
        return False

class TermList(BaseModel):
    terms: List[Term] = Field(description="专业名词列表")

class Classification(BaseModel):
    vertical_classification:str = Field(description="垂直领域一级分类")
    sub_classification:str = Field(description="一级分类下的二级分类")

class QueryRewrite(BaseModel):
    rewrite:str = Field(description="问题改写")

# 1. 软件问题
# 1.1 软件功能
class SoftwareFunction(BaseModel):
    software_name: SoftwareName = Field(description="软件名称,只能从给定的范围中取值")
    function_name: str = Field(description="具体功能名称")
    operation: str = Field(description="用户操作意图(如何使用功能、功能入口、功能使用场景)")
    software_version: Optional[str] = Field(None, description="软件版本")
    operation_steps: Optional[str] = Field(None, description="操作步骤描述")
    
    def check_required_slots(self) -> Tuple[bool, Dict[str, str]]:
        """检查必填槽位是否都存在"""
        missing_slots = {}
        if not self.software_name:
            missing_slots["software_name"] = f"{SoftwareFunction.model_fields['software_name'].description}，可选值：{', '.join([name.value for name in SoftwareName if name not in [SoftwareName.UNKNOWN, SoftwareName.ALIASES]])}"
        if not self.function_name:
            missing_slots["function_name"] = SoftwareFunction.model_fields["function_name"].description
        if not self.operation:
            missing_slots["operation"] = SoftwareFunction.model_fields["operation"].description
        return len(missing_slots) == 0, missing_slots

# 1.2 故障排查
class TroubleShooting(BaseModel):
    software_name: SoftwareName = Field(description="软件名称,只能从给定的范围中取值")
    function_name: str = Field(description="具体功能名称/操作描述")
    error_message: str = Field(description="报错信息/异常现象")
    software_version: Optional[str] = Field(None, description="软件版本")
    os_version: Optional[str] = Field(None, description="操作系统及版本")
    reproduction_steps: Optional[str] = Field(None, description="故障重现步骤")
    
    def check_required_slots(self) -> Tuple[bool, Dict[str, str]]:
        """检查必填槽位是否都存在"""
        missing_slots = {}
        if not self.software_name:
            missing_slots["software_name"] = f"{TroubleShooting.model_fields['software_name'].description}，可选值：{', '.join([name.value for name in SoftwareName if name not in [SoftwareName.UNKNOWN, SoftwareName.ALIASES]])}"
        if not self.function_name:
            missing_slots["function_name"] = TroubleShooting.model_fields["function_name"].description
        if not self.error_message:
            missing_slots["error_message"] = TroubleShooting.model_fields["error_message"].description
        return len(missing_slots) == 0, missing_slots

# 2. 业务问题
# 2.1 专业咨询
class ProfessionalConsulting(BaseModel):
    scene_subject: str = Field(description="场景主体")
    business_scene: str = Field(description="业务场景描述")
    software_name: Optional[SoftwareName] = Field(None, description="软件名称")
    
    def check_required_slots(self) -> Tuple[bool, Dict[str, str]]:
        """检查必填槽位是否都存在"""
        missing_slots = {}
        if not self.scene_subject:
            missing_slots["scene_subject"] = ProfessionalConsulting.model_fields["scene_subject"].description
        if not self.business_scene:
            missing_slots["business_scene"] = ProfessionalConsulting.model_fields["business_scene"].description
        return len(missing_slots) == 0, missing_slots

# 2.2 数据问题
class DataProblem(BaseModel):
    expense_type: str = Field(description="费用类型")
    operation_purpose: str = Field(description="操作目的")
    software_name: Optional[SoftwareName] = Field(None, description="软件名称")
    project_type: Optional[str] = Field(None, description="工程类型")
    
    def check_required_slots(self) -> Tuple[bool, Dict[str, str]]:
        """检查必填槽位是否都存在"""
        missing_slots = {}
        if not self.expense_type:
            missing_slots["expense_type"] = DataProblem.model_fields["expense_type"].description
        if not self.operation_purpose:
            missing_slots["operation_purpose"] = DataProblem.model_fields["operation_purpose"].description
        return len(missing_slots) == 0, missing_slots

# 3. 安装下载注册
# 3.1 后缀名咨询
class FileExtensionConsulting(BaseModel):
    file_extension: str = Field(description="文件后缀名")
    operation_purpose: str = Field(description="操作目的")
    file_source: Optional[str] = Field(None, description="文件来源场景")
    related_software: Optional[str] = Field(None, description="相关软件名称")
    
    def check_required_slots(self) -> Tuple[bool, Dict[str, str]]:
        """检查必填槽位是否都存在"""
        missing_slots = {}
        if not self.file_extension:
            missing_slots["file_extension"] = FileExtensionConsulting.model_fields["file_extension"].description
        if not self.operation_purpose:
            missing_slots["operation_purpose"] = FileExtensionConsulting.model_fields["operation_purpose"].description
        return len(missing_slots) == 0, missing_slots

# 3.2 软件锁类
class SoftwareLock(BaseModel):
    lock_type: str = Field(description="锁类型")
    operation_purpose: str = Field(description="操作目的")
    lock_number: Optional[str] = Field(None, description="软件锁编号/注册号")
    
    def check_required_slots(self) -> Tuple[bool, Dict[str, str]]:
        """检查必填槽位是否都存在"""
        missing_slots = {}
        if not self.lock_type:
            missing_slots["lock_type"] = SoftwareLock.model_fields["lock_type"].description
        if not self.operation_purpose:
            missing_slots["operation_purpose"] = SoftwareLock.model_fields["operation_purpose"].description
        return len(missing_slots) == 0, missing_slots

# 3.3 安装下载类
class InstallationDownload(BaseModel):

    software_name: str = Field(description="软件/插件名称，与file_name二选一")
    file_name: str = Field(description="文件名，与software_name二选一")
    operation_stage: str = Field(description="操作阶段")
    os_version: Optional[str] = Field(None, description="操作系统版本")
    package_source: Optional[str] = Field(None, description="安装包来源/版本号")
    
    def check_required_slots(self) -> Tuple[bool, Dict[str, str]]:
        """检查必填槽位是否都存在"""
        missing_slots = {}
        if not self.software_name and not self.file_name:
            missing_slots["software_name"] = f"{InstallationDownload.model_fields['software_name'].description},"
            f"可选值：{', '.join([name.value for name in SoftwareName if name not in [SoftwareName.UNKNOWN, SoftwareName.ALIASES]])}"
            missing_slots["file_name"] = InstallationDownload.model_fields["file_name"].description
        if not self.operation_stage:
            missing_slots["operation_stage"] = InstallationDownload.model_fields["operation_stage"].description
        return len(missing_slots) == 0, missing_slots

# 3.4 问题排查类
class ProblemDiagnosis(BaseModel):
    error_message: str = Field(description="报错信息/异常现象")
    software_name: Optional[SoftwareName] = Field(None, description="软件名称,只能从给定的范围中取值")
    os_version: Optional[str] = Field(None, description="操作系统版本")
    
    def check_required_slots(self) -> Tuple[bool, Dict[str, str]]:
        """检查必填槽位是否都存在"""
        missing_slots = {}
        if not self.error_message:
            missing_slots["error_message"] = ProblemDiagnosis.model_fields["error_message"].description
        return len(missing_slots) == 0, missing_slots