BIDI_EXTRACTION_AI/BiddingKG/dl/template_extract/templates/XX7523/template1.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:
        """
        检查是否符合模板调用条件：匹配data.json的核心表头（按实际位置）
        :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:]

            # 定义核心表头的泛化正则匹配规则（按data.json实际表头位置）
            # 位置0:收款方 位置1:商品名称 位置2:规格 位置3:最小单位 位置4:单价 位置5:数量 位置6:总价
            header_patterns = [
                r'收款方',  # 第0列：收款方
                r'商品名称|产品名称|货品名称|标的名称',  # 第1列：商品名称
                r'规格|型号|规格型号|技术参数',        # 第2列：规格
                r'最小单位|单位|数量单位|计量单元',    # 第3列：最小单位
                r'单价|单价金额|单价（元）|单价（万元）',# 第4列：单价
                r'数量|采购数量|成交数量|订购数量',      # 第5列：数量
                r'总价|总金额|合计金额|成交总价'        # 第6列：总价
            ]

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

            # 按表头位置提取字段（严格对应data.json的列）
            payee = str(row[0]).strip()        # 收款方（中标人）
            product_name = str(row[1]).strip() # 商品名称
            spec = str(row[2]).strip()         # 规格
            unit = str(row[3]).strip()         # 最小单位
            unit_price = str(row[4]).strip()   # 单价
            quantity = str(row[5]).strip()     # 数量
            total_price = str(row[6]).strip()  # 总价

            # ========== 提取产品信息（必须包含产品+至少一个其他要素） ==========
            if product_name:  # 产品名称为必填
                # 补充单价单位（优先从表头提取，无则默认加"元"）
                price_header = str(self.table_header[4]).strip()
                unit_match = re.search(r'([万亿美欧日]?元)', price_header, re.IGNORECASE)
                if unit_match and unit_price and not re.search(unit_match.group(1), unit_price, re.IGNORECASE):
                    unit_price += unit_match.group(1)
                elif unit_price and not re.search(r'[万亿美欧日]?元', unit_price, re.IGNORECASE):
                    unit_price += "元"

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

                # 构建产品信息字典
                product_info = {
                    "产品": product_name,
                    "品目": "",
                    "品牌": "",
                    "规格": spec,
                    "数量": quantity,
                    "单位": unit,
                    "单价": unit_price,
                    "总价": total_price
                }

                # 过滤空值字段
                product_info_filtered = {k: v for k, v in product_info.items() if v}
                # 产品信息必须包含产品+至少一个其他要素
                if len(product_info_filtered) > 1:
                    extract_result["产品信息"].append(product_info_filtered)

        # ========== 按规则处理其他字段 ==========
        # 招标信息：无招标人/预算，置空（符合要求）
        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("当前数据不满足模板调用条件")