feat: loader prompt & simplify YAML loading and update data loader specifications (#736)

* refactor: Simplify YAML loading and update data loader specifications

* docs: Add comment explaining FunctionLoader usage in tpl.py

* lint

* lint

rdagent/components/coder/data_science/raw_data_loader/prompts.yaml

@@ -51,33 +51,12 @@ spec:
           - Specify the data source location (`/kaggle/input/`).
           - Clearly define the structure and type of the output.
           - Inferred data shape to each input and output data variables. To uncertain dimension, use -1.
-
-      2. Precautions for Data Loading and Preprocessing:
-        - File Handling:
-          - Handle file encoding and delimiters appropriately.
-          - Combine or process multiple files if necessary.
-          - Avoid using the sample submission file to infer test indices. If a dedicated test index file is available, use that. If not, use the order in the test file as the test index.
-        - Data Preprocessing:
-          - Convert data types correctly (e.g., numeric, categorical, date parsing).
-          - Handle missing values appropriately.
-          - Optimize memory usage for large datasets using techniques like downcasting or reading data in chunks if necessary.
-        - Domain-Specific Handling: 
-          - Apply competition-specific preprocessing steps as needed (e.g., text tokenization, image resizing).
-
-      3. Code Standards:
-        - DO NOT use progress bars (e.g., `tqdm`).
-        - DO NOT use the sample submission file to extract test index information.
-        - DO NOT exclude features inadvertently during this process.
-
-      4. Notes:
+      2. Notes:
         - Update `DT` (data type) based on the specific competition dataset. This can include `pd.DataFrame`, `np.array`, `torch.Tensor`, etc.
-        - Never use sample submission as the test index, as it may not be the same as the test data. Use the test index file or test data source to get the test index.
         - Only set the DT of variables without inferring the shape of these variables since you don't know the shape of the data.
-      
-      5. Exploratory Data Analysis (EDA) [Required]:
-        - Before returning the data, you should always add an EDA part describing the data to help the following steps understand the data better.
-        - The EDA part should be drafted in plain text with certain format schema with no more than ten thousand characters.
-        - An evaluation agent will help to check whether the EDA part is added correctly.
+
+      Responsibilities and notes of an implemented data loader that aligns with the generated specification.
+      {% include "scenarios.data_science.share:component_spec.DataLoadSpec" %}
 
       {% if latest_spec %}
       6. Former Specification:

rdagent/scenarios/data_science/share.yaml

@@ -64,7 +64,7 @@ describe: # some template to describe some object
 
 component_description:
   DataLoadSpec: |-
-    Loads and preprocesses competition data, ensuring proper data types, handling missing values, and providing an exploratory data analysis summary.
+    Loads raw competition data, ensuring proper data types, and providing an exploratory data analysis summary.
   FeatureEng: |-
     Transforms raw data into meaningful features while maintaining shape consistency, avoiding data leakage, and optimizing for model performance.
     It should be model-agnostic (data transformations/augmentations that apply only to specific model frameworks should not be included here).
@@ -89,22 +89,27 @@ component_spec:
       - Handle file encoding and delimiters appropriately.
       - Combine or process multiple files if necessary.
       - Avoid using the sample submission file to infer test indices. If a dedicated test index file is available, use that. If not, use the order in the test file as the test index.
+      - Ensure you load the actual data from the files, not just the filenames or paths. Do not postpone data loading to later steps.
 
     2. Data Preprocessing:
       - Convert data types correctly (e.g., numeric, categorical, date parsing).
-      - Handle missing values appropriately.
       - Optimize memory usage for large datasets using techniques like downcasting or reading data in chunks if necessary.
+      - Domain-Specific Handling: 
+        - Apply competition-specific preprocessing steps as needed (e.g., text tokenization, image resizing).
 
     3. Code Standards:
       - DO NOT use progress bars (e.g., `tqdm`).
       - DO NOT use the sample submission file to extract test index information.
       - DO NOT exclude features inadvertently during this process.
-    
+
     4. Exploratory Data Analysis (EDA) [Required]:
       - Before returning the data, you should always add an EDA part describing the data to help the following steps understand the data better.
       - The EDA part should be drafted in plain text with certain format schema with no more than ten thousand characters.
       - An evaluation agent will help to check whether the EDA part is added correctly.
 
+    5. NOTES
+      - Never use sample submission as the test index, as it may not be the same as the test data. Use the test index file or test data source to get the test index.
+
   FeatureEng: |-
     1. Well handle the shape of the data
       - The sample size of the train data and the test data should be the same in all scenarios.

rdagent/utils/agent/tpl.py

@@ -9,15 +9,52 @@ from pathlib import Path
 from typing import Any
 
 import yaml
-from jinja2 import Environment, StrictUndefined
+from jinja2 import Environment, FunctionLoader, StrictUndefined
 
-from rdagent.core.utils import SingletonBaseClass
 from rdagent.log import rdagent_logger as logger
 
 DIRNAME = Path(__file__).absolute().resolve().parent
 PROJ_PATH = DIRNAME.parent.parent
 
 
+def get_caller_dir(upshift: int = 0) -> Path:
+    # Inspect the calling stack to get the caller's directory
+    stack = inspect.stack()
+    caller_frame = stack[1 + upshift]
+    caller_module = inspect.getmodule(caller_frame[0])
+    if caller_module and caller_module.__file__:
+        caller_dir = Path(caller_module.__file__).parent
+    else:
+        caller_dir = DIRNAME
+    return caller_dir
+
+
+def load_yaml_content(uri: str, caller_dir: Path | None = None) -> Any:
+    """
+    Please refer to RDAT.__init__ file
+    """
+    if caller_dir is None:
+        caller_dir = get_caller_dir(upshift=1)
+    # Parse the URI
+    path_part, yaml_path = uri.split(":")
+    yaml_keys = yaml_path.split(".")
+
+    if path_part.startswith("."):
+        yaml_file_path = caller_dir / f"{path_part[1:].replace('.', '/')}.yaml"
+    else:
+        yaml_file_path = (PROJ_PATH / path_part.replace(".", "/")).with_suffix(".yaml")
+
+    # Load the YAML file
+    with open(yaml_file_path, "r") as file:
+        yaml_content = yaml.safe_load(file)
+
+    # Traverse the YAML content to get the desired template
+    for key in yaml_keys:
+        yaml_content = yaml_content[key]
+
+    return yaml_content
+
+
 # class T(SingletonBaseClass): TODO: singleton does not support args now.
 class RDAT:
     """
@@ -39,44 +76,27 @@ class RDAT:
             the loaded content will be saved in `self.template`
         """
         self.uri = uri
-        # Inspect the calling stack to get the caller's directory
-        stack = inspect.stack()
-        caller_frame = stack[1]
-        caller_module = inspect.getmodule(caller_frame[0])
-        if caller_module and caller_module.__file__:
-            caller_dir = Path(caller_module.__file__).parent
-        else:
-            caller_dir = DIRNAME
-
-        # Parse the URI
-        path_part, yaml_path = uri.split(":")
-        yaml_keys = yaml_path.split(".")
-
-        if path_part.startswith("."):
-            yaml_file_path = caller_dir / f"{path_part[1:].replace('.', '/')}.yaml"
+        caller_dir = get_caller_dir(1)
+        if uri.startswith("."):
             try:
                 # modify the uri to a raltive path to the project for easier finding prompts.yaml
                 self.uri = f"{str(caller_dir.resolve().relative_to(PROJ_PATH)).replace('/', '.')}{uri}"
             except ValueError:
                 pass
-        else:
-            yaml_file_path = (PROJ_PATH / path_part.replace(".", "/")).with_suffix(".yaml")
-
-        # Load the YAML file
-        with open(yaml_file_path, "r") as file:
-            yaml_content = yaml.safe_load(file)
-
-        # Traverse the YAML content to get the desired template
-        for key in yaml_keys:
-            yaml_content = yaml_content[key]
-
-        self.template = yaml_content
+        self.template = load_yaml_content(uri, caller_dir=caller_dir)
 
     def r(self, **context: Any) -> str:
         """
         Render the template with the given context.
         """
-        rendered = Environment(undefined=StrictUndefined).from_string(self.template).render(**context).strip("\n")
+        # loader=FunctionLoader(load_yaml_content) is for supporting grammar like below.
+        # `{% include "scenarios.data_science.share:component_spec.DataLoadSpec" %}`
+        rendered = (
+            Environment(undefined=StrictUndefined, loader=FunctionLoader(load_yaml_content))
+            .from_string(self.template)
+            .render(**context)
+            .strip("\n")
+        )
         while "\n\n\n" in rendered:
             rendered = rendered.replace("\n\n\n", "\n\n")
         rendered = "\n".join(line for line in rendered.splitlines() if line.strip())

test/utils/test_agent_infra.py

@@ -24,6 +24,12 @@ class TestAgentInfra(unittest.TestCase):
 
         print(code)
 
+    def test_include(self):
+        parent = T("components.coder.data_science.raw_data_loader.prompts:spec.user.data_loader").r(latest_spec=None)
+        child = T("scenarios.data_science.share:component_spec.DataLoadSpec").r()
+        assert child in parent
+        print(parent)
+
 
 if __name__ == "__main__":
     unittest.main()