BIDI_EXTRACTION_AI/BiddingKG/dl/template_extract/templates/shared_template/template2.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:
            return False
        tables = 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:]

            # 定义核心表头的泛化正则匹配规则（按实际表头位置）
            header_patterns = [
                r'状态|成交状态|中标状态',  # 第0列
                r'供应商名称|成交供应商名称|中标供应商名称|中标人名称',  # 第1列
                r'评审结果|中选结果|中标结果',  # 第2列
                r'公告日期|发布日期|公示日期',  # 第3列
                r'成交金额|中标金额|成交价|中标价|金额',  # 第4列
                r'优惠率|折扣率|让利幅度'  # 第5列
            ]

            # 验证表头匹配（按位置校验）
            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  # 跳过长度不匹配的行

            # 按表头位置提取字段
            supplier_name = str(row[1]).strip()  # 供应商名称（中标人）
            deal_amount = str(row[4]).strip()    # 成交金额
            review_result = str(row[2]).strip()  # 评审结果

            # 跳过无中标人的行（中标信息必须包含中标人）
            if not supplier_name or review_result not in ['中选', '中标']:
                continue

            # 补充金额单位（优先从表头提取，无则默认加"元"）
            amount_header = str(self.table_header[4]).strip()
            unit_match = re.search(r'([万亿美欧日]?元)', amount_header, re.IGNORECASE)
            if unit_match and not re.search(unit_match.group(1), deal_amount, re.IGNORECASE):
                deal_amount += unit_match.group(1)
            elif deal_amount and not re.search(r'[万亿美欧日]?元', deal_amount, re.IGNORECASE):
                deal_amount += "元"

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

            # 过滤空值字段
            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["招标信息"] = []

        # 候选人信息：无候选人+排名，置空
        extract_result["候选人信息"] = []

        # 产品信息：无产品及其他要素，置空
        extract_result["产品信息"] = []

        # 过滤顶层空值字段（项目编号/名称为空则移除）
        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("当前数据不满足模板调用条件")