上传代码

supplement_kg.py

@@ -333,50 +333,47 @@ class ExpenseProcessor:
         return successful_files
 
 
+import re
+
+
 def _determine_project_type(data):
     """
-    根据division字段判断工程类型
+    根据basicData中的"项目类型"或"工程类型"判断工程类型
     :param data: 项目数据
     :return: 'inventory' 表示清单工程，'budget' 表示预算工程
     """
-    # 清单工程关键词
-    inventory_keywords = ["清单", "结算", "招标控制价", "招投标工程", "清单计价"]
-    # 预算工程关键词
-    budget_keywords = ["概预算", "定额", "定额计价", "概算", "概预算工程"]
+    # 项目类型名称映射字典：将各种变体映射到标准类型（预算/清单）
+    PROJECT_TYPE_MAPPING = {
+        "概预算工程": "预算",
+        "初步设计概算": "预算",
+        "可行性研究投资估算": "预算",
+        "施工图预算": "预算",
+        "配网定额计价": "预算",
+        "招标控制价": "清单",
+        "投标报价": "清单",
+        "招投标工程": "清单",
+        "配网清单招投标计价": "清单",
+    }
 
-    # 尝试从数据中获取division字段
-    division = None
-    if "division" in data:
-        division = data["division"]
-        parts = division.split("-")
+    # 获取 basicData
+    basic_data = data.get("basicData") or {}
 
-    # 如果找到division字段
-    if division:
-        # 去掉"主网-"前缀
-        if len(parts) == 2:
-            division_type = parts[1].strip()
+    # 尝试获取 "项目类型"，若不存在则尝试获取 "工程类型"
+    engineering_type = basic_data.get("项目类型") or basic_data.get("工程类型") or basic_data.get("工程类别")
+
+    if engineering_type:
+        # 去除前后空格
+        engineering_type = engineering_type.strip()
+        # 查找映射
+        mapped_type = PROJECT_TYPE_MAPPING.get(engineering_type)
+        if mapped_type == "预算":
+            print(f"根据项目类型 '{engineering_type}' 判断为预算工程")
+            return "budget"
+        elif mapped_type == "清单":
+            print(f"根据项目类型 '{engineering_type}' 判断为清单工程")
+            return "inventory"
         else:
-            division_type = parts[2].strip()
-
-        # 判断是否为清单工程
-        for keyword in inventory_keywords:
-            if keyword in division_type:
-                print(f"根据division字段 '{division}' 判断为清单工程")
-                return "inventory"
-
-        # 判断是否为预算工程
-        for keyword in budget_keywords:
-            if keyword in division_type:
-                print(f"根据division字段 '{division}' 判断为预算工程")
-                return "budget"
-
-    # 如果无法通过division字段判断，则尝试通过数据结构判断
-    is_inventory_project = False
-    for key in data.keys():
-        if re.search(r"[0-9A-F]{8}-[0-9A-F]{4}-[0-9A-F]{4}-[0-9A-F]{4}-[0-9A-F]{12}", key, re.IGNORECASE):
-            is_inventory_project = True
-            print("通过数据结构判断为清单工程")
-            break
+            print(f"项目类型 '{engineering_type}' 未在映射中定义，跳过")
 
     return "inventory" if is_inventory_project else "budget"
 
@@ -388,8 +385,8 @@ def costsummary_upwards(
 
 
 if __name__ == "__main__":
-    input_directory = "project2json/outputs/json"
-    output_directory = "project2json/outputs/merged"
+    input_directory = "data/output/json"
+    output_directory = "data/output/merged"
     # 自动判断工程类型
     result = costsummary_upwards(input_directory, output_directory)
     if result: