BIDI_EXTRACTION_AI/BiddingKG/dl/template_extract/templates/XX0182/template6.py

# coding:utf8

import re
from typing import Dict, Any
from BiddingKG.dl.template_extract.abstract_template import AbstractTemplate
import os
import json


class ProductInfoTemplate(AbstractTemplate):
    """
    产品信息提取模板（适配data.json中的第二个表格结构）
    从第二个表格提取产品相关信息
    """

    def __init__(self):
        super().__init__(
            template_id=os.path.abspath(__file__),
            priority=3
        )
        self.target_table = None  # 第二个表格，存储产品信息

    def check_call_timing(self, preprocessed_data: Dict[str, Any]) -> bool:
        """检查是否符合模板调用条件"""
        if "表格列表" not in preprocessed_data:
            return False

        tables = preprocessed_data["表格列表"]
        if len(tables) < 1:
            print("表格数量不足，至少需要2个表格")
            return False

        try:
            self.target_table = tables[0]  # 只处理第二个表格

            # 验证表格关键表头（按data.json第二个表格实际结构泛化匹配）
            header_patterns = [
                (0, 0, r'采购品目|品目'),  # 第1行第1列
                (0, 2, r'名称|产品名称'),  # 第1行第3列
                (1, 0, r'单位|计量单位'),  # 第2行第1列
                (1, 2, r'数量'),  # 第2行第3列
                (2, 0, r'单价'),  # 第3行第1列
                (2, 2, r'总价|总金额'),  # 第3行第3列
            ]

            # 验证表头
            for row, col, pattern in header_patterns:
                if row >= len(self.target_table) or col >= len(self.target_table[row]):
                    return False
                header_text = self.target_table[row][col].strip()
                if not re.search(pattern, header_text, re.IGNORECASE):
                    print(f"表头不匹配: {header_text} 不符合 {pattern}")
                    return False

            return True
        except Exception as e:
            print(f"表格验证错误: {e}")
            return False

    def extract(self, preprocessed_data: Dict[str, Any]) -> Dict[str, Any]:
        """提取表格中的产品信息"""
        result = {"产品信息": []}

        # 提取产品信息
        product_info = self._extract_product_info()
        if product_info:
            result["产品信息"].append(product_info)

        return result

    def _get_value(self, table: list, row: int, col: int) -> str:
        """获取表格指定位置的值"""
        if row < len(table) and col < len(table[row]):
            return table[row][col].strip()
        return ""

    def _format_amount(self, value: str, header_row: int, header_col: int) -> str:
        """格式化金额，补充单位"""
        if not value:
            return ""
        # 检查值中是否已包含单位
        if re.search(r'[万亿美欧日]?元|￥|人民币', value):
            return value
        # 从表头获取单位补充
        header_text = self._get_value(self.target_table, header_row, header_col)
        if re.search(r'[万亿美欧日]?元|￥|人民币', header_text):
            unit = re.search(r'[万亿美欧日]?元|￥|人民币', header_text).group(0)
            return f"{value}{unit}"
        return value

    def _extract_product_info(self) -> Dict[str, Any]:
        """提取产品详细信息"""
        product = {
            "产品": self._get_value(self.target_table, 0, 3),  # 产品名称
            "品目": self._get_value(self.target_table, 0, 1),  # 采购品目
            "数量": self._get_value(self.target_table, 1, 3),  # 数量
            "单位": self._get_value(self.target_table, 1, 1),  # 计量单位
            "单价": self._format_amount(self._get_value(self.target_table, 2, 1), 2, 0),  # 单价（补充单位）
            "总价": self._format_amount(self._get_value(self.target_table, 2, 3), 2, 2),  # 总价（补充单位）
            # 品牌和规格在当前表格中不存在，不提取
        }
        # 过滤空值字段
        return {k: v for k, v in product.items() if v}


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

    template = ProductInfoTemplate()
    if template.check_call_timing(preprocessed_data):
        result = template.extract(preprocessed_data)
        print("提取结果：", json.dumps(result, ensure_ascii=False, indent=2))
    else:
        print("当前数据不满足模板调用条件")