BIDI_EXTRACTION_AI/BiddingKG/dl/template_extract/templates/DX006116/template3.py

# coding:utf8
import re
from typing import Dict, Any
from BiddingKG.dl.template_extract.abstract_template import AbstractTemplate
import os
import json


class CustomTransactionTemplate(AbstractTemplate):
    """
    适配data.json表格结构的成交候选人信息提取模板
    提取成交候选人信息、过滤空值，严格遵循返回格式要求
    """

    def __init__(self):
        super().__init__(
            template_id=os.path.abspath(__file__),
            priority=3
        )
        self.table_header = None  # 表格表头
        self.table_data = None    # 表格数据行

    def check_call_timing(self, preprocessed_data: Dict[str, Any]) -> bool:
        """
        检查是否符合模板调用条件：匹配表格结构和核心表头
        :param preprocessed_data: 预处理后的输入数据（包含表格列表）
        :return: 是否符合调用条件
        """
        # 检查表格列表是否存在且有效
        if "表格列表" not in preprocessed_data or '表格前文' not in preprocessed_data:
            return False
        tables = preprocessed_data["表格列表"]
        tables_prev_text = preprocessed_data['表格前文']
        if len(tables) < 1 or not isinstance(tables[0], list) or len(tables[0]) < 2:
            return False

        try:
            # 提取第一个表格的表头和数据
            target_table = tables[0]
            self.table_header = target_table[0]
            self.table_data = target_table[1:]
            prev_text = re.sub('\s', '', tables_prev_text[0])

            if len(self.table_header) != 3:
                # print(f"表头列数异常：{len(self.table_header)}列，预期3列")
                return False

            if re.search('\d、(中标|成交|中选)(候选)?人$', prev_text) == None:
                return False

            # 定义核心表头的泛化正则匹配规则（按实际表头位置）
            header_patterns = [
                r'包件号',  # 第0列
                r'包件名称',  # 第1列
                r'成交候选(人|单位)|中标候选(人|单位)',  # 第2列，泛化匹配成交/中标候选人
            ]

            # 验证表头匹配（按位置校验，忽略大小写）
            for idx, pattern in enumerate(header_patterns):
                if idx >= len(self.table_header):
                    return False
                header_str = str(self.table_header[idx]).strip()
                if not re.fullmatch(pattern, header_str, re.IGNORECASE):
                    # print(f"表头位置{idx}匹配失败: 预期{pattern}，实际{header_str}")
                    return False

            return True
        except Exception as e:
            print(f"校验表格失败: {e}")
            return False

    def extract(self, preprocessed_data: Dict[str, Any]) -> Dict[str, Any]:
        """
        提取表格数据，按指定格式返回（过滤空值、校验字段有效性）
        :param preprocessed_data: 预处理后的输入数据
        :return: 提取结果（符合指定JSON格式）
        """
        # 初始化返回结果结构
        extract_result = {
            "项目编号": "",
            "项目名称": "",
            "招标信息": [],
            "中标信息": [],
            "候选人信息": [],
            "产品信息": []
        }

        # 提取候选人信息（核心逻辑）
        for row in self.table_data:
            if len(row) != len(self.table_header):
                continue  # 跳过长度不匹配的行

            # 按表头位置提取字段
            section_num = str(row[0]).strip()  # 包件号
            package_num = str(row[1]).strip()  # 包件名称
            candidate_name = str(row[2]).strip()  # 成交候选人名称

            # 构建中标信息字典
            bid_info = {
                "标的": "",
                "标包": package_num,
                "包号": section_num,
                "中标人": candidate_name,
                "中标价": '',
                "联系人": "",
                "电话": "",
                "服务时间": "",
                "地址": ""
            }

            # 过滤空值字段
            bid_info_filtered = {k: v for k, v in bid_info.items() if v}
            if bid_info_filtered:  # 确保非空才添加
                extract_result["中标信息"].append(bid_info_filtered)


        # 过滤顶层空值字段（项目编号/名称为空则移除）
        extract_result = {k: v for k, v in extract_result.items() if v or isinstance(v, list)}

        return extract_result


# 使用示例
if __name__ == "__main__":
    # 读取data.json数据
    with open("data.json", "r", encoding="utf8") as f:
        preprocessed_data = json.load(f)

    # 初始化模板并执行提取
    template = CustomTransactionTemplate()
    if template.check_call_timing(preprocessed_data):
        result = template.extract(preprocessed_data)
        # 输出格式化JSON结果
        print("提取结果：")
        print(json.dumps(result, ensure_ascii=False, indent=4))
    else:
        print("当前数据不满足模板调用条件")