Merge pull request #415 from shijinpjlab/dev_0601

feat: 结果保存优化：数据库字段序列化、ExecutorResultSaveArgs新增参数

Author: shijinpjlab

dingo/config/input_args.py

@@ -77,6 +77,9 @@ class ExecutorResultSaveArgs(BaseModel):
     all_labels: bool = False
     raw: bool = False
     merge: bool = False  # 如果为True，所有数据写入同一个jsonl文件，不分文件夹
+    limit: Optional[int] = None  # 每个输出文件最多写入条数，None表示不限制
+    field_list: Optional[List[str]] = None  # 仅保存指定字段；若均不存在则报错
+    full_field_sample_count: int = 0  # 保留完整字段样本条数，0表示关闭
 
 
 class ExecutorArgs(BaseModel):

dingo/exec/local.py

@@ -5,7 +5,9 @@
 import os
 import time
 import uuid
-from typing import Generator, List, Optional
+from datetime import date, datetime
+from decimal import Decimal
+from typing import Dict, Generator, List, Optional
 
 from tqdm import tqdm
 
@@ -25,6 +27,8 @@ def __init__(self, input_args: InputArgs):
         self.input_args: InputArgs = input_args
         self.llm: Optional[BaseLLM] = None
         self.summary: SummaryModel = SummaryModel()
+        self._full_field_written_count: int = 0
+        self._file_written_count: Dict[str, int] = {}
 
     def load_data(self) -> Generator[Data, None, None]:
         """
@@ -252,6 +256,28 @@ def summarize(self, summary: SummaryModel) -> SummaryModel:
         new_summary.finish_time = time.strftime("%Y%m%d_%H%M%S", time.localtime())
         return new_summary
 
+    @staticmethod
+    def _json_default(value):
+        if isinstance(value, Decimal):
+            return float(value)
+        if isinstance(value, (datetime, date)):
+            return value.isoformat()
+        return str(value)
+
+    @classmethod
+    def _json_dumps(cls, value: dict) -> str:
+        return json.dumps(value, ensure_ascii=False, default=cls._json_default)
+
+    def _resolve_field_list(self, input_args: InputArgs) -> Optional[List[str]]:
+        if input_args.executor.result_save.full_field_sample_count <= 0:
+            return input_args.executor.result_save.field_list
+
+        if self._full_field_written_count < input_args.executor.result_save.full_field_sample_count:
+            self._full_field_written_count += 1
+            return None
+
+        return input_args.executor.result_save.field_list
+
     def write_single_data(
         self, path: str, input_args: InputArgs, result_info: ResultInfo
     ):
@@ -261,20 +287,20 @@ def write_single_data(
         # 如果启用 merge 模式，将所有数据写入同一个文件
         if input_args.executor.result_save.merge:
             f_n = os.path.join(path, "all_results.jsonl")
+            if not self._should_write_to_file(input_args, f_n):
+                return
+            str_json = self._build_output_json(input_args, result_info)
             with open(f_n, "a", encoding="utf-8") as f:
-                # if input_args.executor.result_save.raw:
-                #     str_json = json.dumps(result_info.to_raw_dict(), ensure_ascii=False)
-                # else:
-                #     str_json = json.dumps(result_info.to_dict(), ensure_ascii=False)
-                str_json = json.dumps(result_info.to_raw_dict(), ensure_ascii=False)
                 f.write(str_json + "\n")
+            self._file_written_count[f_n] = self._file_written_count.get(f_n, 0) + 1
             return
 
         if not input_args.executor.result_save.good and not result_info.eval_status:
             return
 
         # 用集合记录已经写过的(字段名, label名)组合，避免重复写入
         written_labels = set()
+        str_json: Optional[str] = None
 
         # 遍历 eval_details 的第一层（字段名组合），第二层是List[EvalDetail]
         for field_name, eval_detail_list in result_info.eval_details.items():
@@ -314,18 +340,39 @@ def write_single_data(
                         # 没有点分割，直接在字段文件夹下创建文件
                         f_n = os.path.join(field_dir, parts[0] + ".jsonl")
 
+                    if not self._should_write_to_file(input_args, f_n):
+                        continue
+                    if str_json is None:
+                        str_json = self._build_output_json(input_args, result_info)
                     with open(f_n, "a", encoding="utf-8") as f:
-                        if input_args.executor.result_save.raw:
-                            str_json = json.dumps(result_info.to_raw_dict(), ensure_ascii=False)
-                        else:
-                            str_json = json.dumps(result_info.to_dict(), ensure_ascii=False)
                         f.write(str_json + "\n")
+                    self._file_written_count[f_n] = self._file_written_count.get(f_n, 0) + 1
+
+    def _should_write_to_file(self, input_args: InputArgs, file_path: str) -> bool:
+        limit = input_args.executor.result_save.limit
+        if limit is None:
+            return True
+        return self._file_written_count.get(file_path, 0) < limit
+
+    def _build_output_json(self, input_args: InputArgs, result_info: ResultInfo) -> str:
+        field_list = self._resolve_field_list(input_args)
+        if input_args.executor.result_save.raw:
+            output_data = result_info.to_raw_dict(field_list=field_list)
+        else:
+            output_data = result_info.to_dict(field_list=field_list)
+        return self._json_dumps(output_data)
 
     def write_summary(self, path: str, input_args: InputArgs, summary: SummaryModel):
         if not input_args.executor.result_save.bad:
             return
         with open(path + "/summary.json", "w", encoding="utf-8") as f:
-            json.dump(summary.to_dict(), f, indent=4, ensure_ascii=False)
+            json.dump(
+                summary.to_dict(),
+                f,
+                indent=4,
+                ensure_ascii=False,
+                default=self._json_default,
+            )
 
     def get_summary(self):
         return self.summary

dingo/io/output/result_info.py

@@ -1,4 +1,9 @@
-from typing import Dict, List
+import ast
+import json
+from collections.abc import Mapping
+from datetime import date, datetime
+from decimal import Decimal
+from typing import Any, Dict, List, Optional
 
 from pydantic import BaseModel
 
@@ -11,38 +16,96 @@ class ResultInfo(BaseModel):
     eval_status: bool = False
     eval_details: Dict[str, List[EvalDetail]] = {}
 
-    def to_dict(self):
+    @staticmethod
+    def _apply_field_filter(output_data: Dict[str, Any], field_list: Optional[List[str]]) -> Dict[str, Any]:
+        if field_list is None:
+            return output_data
+
+        if len(field_list) == 0:
+            return {}
+
+        missing_fields = [field for field in field_list if field not in output_data]
+        if missing_fields:
+            sample_keys = list(output_data.keys())[:20]
+            raise ValueError(
+                f"result_save.field_list 中字段不存在: {missing_fields}。"
+                f"可用字段示例: {sample_keys}"
+            )
+
+        return {field: output_data[field] for field in field_list}
+
+    @classmethod
+    def _parse_container_string(cls, value: str) -> Any:
+        text = value.strip()
+        if len(text) < 2:
+            return value
+        if not (
+            (text.startswith("[") and text.endswith("]"))
+            or (text.startswith("{") and text.endswith("}"))
+        ):
+            return value
+        try:
+            parsed = json.loads(text)
+        except json.JSONDecodeError:
+            try:
+                parsed = ast.literal_eval(text)
+            except (ValueError, SyntaxError):
+                return value
+        if isinstance(parsed, (dict, list, tuple, set)):
+            return cls._normalize_value(parsed)
+        return value
+
+    @classmethod
+    def _normalize_value(cls, value: Any) -> Any:
+        if isinstance(value, (str, int, float, bool)) or value is None:
+            if isinstance(value, str):
+                return cls._parse_container_string(value)
+            return value
+        if isinstance(value, Decimal):
+            return float(value)
+        if isinstance(value, (datetime, date)):
+            return value.isoformat()
+        if isinstance(value, Mapping):
+            return {str(k): cls._normalize_value(v) for k, v in value.items()}
+        if isinstance(value, (list, tuple, set)):
+            return [cls._normalize_value(item) for item in value]
+        return value
+
+    def to_dict(self, field_list: Optional[List[str]] = None):
         """将ResultInfo转换为字典格式
 
         Returns:
             包含所有字段的字典，其中eval_details被转换为嵌套字典结构
         """
-        return {
+        output_data = {
             'dingo_id': self.dingo_id,
-            'raw_data': self.raw_data,
+            'raw_data': self._normalize_value(self.raw_data),
             'eval_status': self.eval_status,
             'eval_details': {
                 k: [model_res.model_dump() for model_res in v]
                 for k, v in self.eval_details.items()
             },
         }
+        return self._apply_field_filter(output_data, field_list)
 
-    def to_raw_dict(self):
+    def to_raw_dict(self, field_list: Optional[List[str]] = None):
         """将ResultInfo合并到raw_data中
 
         Returns:
             包含原始数据和dingo_result的字典
         """
+        merged_raw_data = self._normalize_value(self.raw_data)
+
         def move_conflict_field(field_name: str):
-            if field_name not in self.raw_data:
+            if field_name not in merged_raw_data:
                 return
 
             index = 1
             while True:
                 backup_field = f'{field_name}_old_v{index}'
-                if backup_field not in self.raw_data:
-                    self.raw_data[backup_field] = self.raw_data[field_name]
-                    del self.raw_data[field_name]
+                if backup_field not in merged_raw_data:
+                    merged_raw_data[backup_field] = merged_raw_data[field_name]
+                    del merged_raw_data[field_name]
                     return
                 index += 1
 
@@ -55,6 +118,6 @@ def move_conflict_field(field_name: str):
         }
         move_conflict_field('dingo_id')
         move_conflict_field('dingo_result')
-        self.raw_data['dingo_id'] = self.dingo_id
-        self.raw_data['dingo_result'] = dingo_result
-        return self.raw_data
+        merged_raw_data['dingo_id'] = self.dingo_id
+        merged_raw_data['dingo_result'] = dingo_result
+        return self._apply_field_filter(merged_raw_data, field_list)

dingo/model/model.py

@@ -124,18 +124,28 @@ def load_model(cls):
         if cls.module_loaded:
             return
         this_module_directory = os.path.dirname(os.path.abspath(__file__))
-        # rule auto register
-        for file in os.listdir(os.path.join(this_module_directory, "rule")):
-            path = os.path.join(this_module_directory, "rule", file)
-            if (
-                os.path.isfile(path)
-                and file.endswith(".py")
-                and not file == "__init__.py"
-            ):
-                try:
-                    importlib.import_module("dingo.model.rule." + file.split(".")[0])
-                except ModuleNotFoundError as e:
-                    log.debug(e)
+        # rule auto register - 递归扫描子目录
+        rule_base_dir = os.path.join(this_module_directory, "rule")
+        for root, dirs, files in os.walk(rule_base_dir):
+            dirs[:] = [d for d in dirs if d != "__pycache__"]
+
+            for file in files:
+                if file.endswith(".py") and file != "__init__.py":
+                    rel_path = os.path.relpath(root, rule_base_dir)
+                    if rel_path == ".":
+                        module_name = f"dingo.model.rule.{file[:-3]}"
+                    else:
+                        rel_module = rel_path.replace(os.sep, ".")
+                        module_name = f"dingo.model.rule.{rel_module}.{file[:-3]}"
+
+                    try:
+                        importlib.import_module(module_name)
+                    except ModuleNotFoundError as e:
+                        log.debug(e)
+                    except ImportError as e:
+                        log.debug("=" * 30 + " ImportError " + "=" * 30)
+                        log.debug(f"module {module_name} not imported because: \n{e}")
+                        log.debug("=" * 73)
 
         # llm auto register - 递归扫描子目录
         llm_base_dir = os.path.join(this_module_directory, "llm")

dingo/model/rule/scibase/rule_quanliang.py

@@ -623,7 +623,7 @@ class RuleQuanliangFieldValidation(BaseRule):
         "evaluation_results": "",
     }
 
-    _required_fields = [RequiredField.METADATA]
+    _required_fields = []
     dynamic_config = EvaluatorRuleArgs(key_list=list(FIELD_VALIDATORS.keys()))
 
     @classmethod

test/scripts/exec/test_local.py

@@ -7,6 +7,83 @@
 
 
 class TestLocal:
+    def test_write_single_data_limit_per_file(self, tmp_path):
+        input_data = {
+            "executor": {
+                "result_save": {
+                    "bad": True,
+                    "limit": 1,
+                }
+            },
+            "evaluator": [],
+        }
+        input_args = InputArgs(**input_data)
+        executor = LocalExecutor(input_args)
+
+        result_info = ResultInfo(
+            dingo_id="1",
+            raw_data={"content": "test"},
+            eval_status=True,
+            eval_details={
+                "content": [
+                    EvalDetail(
+                        metric="RuleColonEnd",
+                        status=True,
+                        label=["QUALITY_BAD_EFFECTIVENESS.RuleColonEnd"],
+                        reason=["bad sample"],
+                    )
+                ]
+            },
+        )
+        output_path = str(tmp_path)
+
+        executor.write_single_data(output_path, input_args, result_info)
+        executor.write_single_data(output_path, input_args, result_info)
+
+        target_file = tmp_path / "content" / "QUALITY_BAD_EFFECTIVENESS" / "RuleColonEnd.jsonl"
+        assert target_file.exists()
+        lines = target_file.read_text(encoding="utf-8").strip().splitlines()
+        assert len(lines) == 1
+
+    def test_write_single_data_limit_per_file_merge_mode(self, tmp_path):
+        input_data = {
+            "executor": {
+                "result_save": {
+                    "bad": True,
+                    "merge": True,
+                    "limit": 1,
+                }
+            },
+            "evaluator": [],
+        }
+        input_args = InputArgs(**input_data)
+        executor = LocalExecutor(input_args)
+
+        result_info = ResultInfo(
+            dingo_id="1",
+            raw_data={"content": "test"},
+            eval_status=True,
+            eval_details={
+                "content": [
+                    EvalDetail(
+                        metric="RuleColonEnd",
+                        status=True,
+                        label=["QUALITY_BAD_EFFECTIVENESS.RuleColonEnd"],
+                        reason=["bad sample"],
+                    )
+                ]
+            },
+        )
+        output_path = str(tmp_path)
+
+        executor.write_single_data(output_path, input_args, result_info)
+        executor.write_single_data(output_path, input_args, result_info)
+
+        target_file = tmp_path / "all_results.jsonl"
+        assert target_file.exists()
+        lines = target_file.read_text(encoding="utf-8").strip().splitlines()
+        assert len(lines) == 1
+
     def test_merge_result_info(self):
         existing_list = []
         new_item1 = ResultInfo(