update all (#530)

rdagent/components/coder/data_science/raw_data_loader/eval.py

@@ -58,12 +58,22 @@ class DataLoaderCoSTEEREvaluator(CoSTEEREvaluator):
         implementation.inject_files(**{fname: test_code})
         stdout = implementation.execute(env=de, entry=f"python {fname}")
 
+        if "main.py" in implementation.file_dict:
+            workflow_stdout = implementation.execute(env=de, entry="python main.py")
+        else:
+            workflow_stdout = None
+
         system_prompt = T(".prompts:data_loader_eval.system").r(
             task_desc=target_task.get_task_information(),
             test_code=test_code,
             code=implementation.file_dict["load_data.py"],
+            workflow_stdout=workflow_stdout,
+            workflow_code=implementation.all_codes,
+        )
+        user_prompt = T(".prompts:data_loader_eval.user").r(
+            stdout=stdout,
+            workflow_stdout=workflow_stdout,
         )
-        user_prompt = T(".prompts:data_loader_eval.user").r(stdout=stdout)
 
         resp = APIBackend().build_messages_and_create_chat_completion(user_prompt, system_prompt, json_mode=True)
         return DataLoaderEvalFeedback(**json.loads(resp))

rdagent/components/coder/data_science/raw_data_loader/eval_tests/data_loader_test.txt

@@ -8,18 +8,45 @@ Please make sure the stdout is rich enough to support informative feedback
 """
 
 import pickle
+
+import pandas as pd
 from load_data import load_data
 
 X, y, X_test, test_ids = load_data()
 
+
 def get_length(data):
     return len(data) if isinstance(data, list) else data.shape[0]
 
-assert get_length(X_test) == get_length(test_ids), (
-    f"Mismatch in length of test images and test IDs: X_test ({get_length(X_test)}) and test_ids ({get_length(test_ids)})"
-)
-assert get_length(X) == get_length(y), (
-    f"Mismatch in length of training images and labels: X ({get_length(X)}) and y ({get_length(y)})"
-)
+
+def get_width(data):
+    return 1 if isinstance(data, list) else data.shape[1:]
+
+
+def get_column_list(data):
+    return data.columns.tolist() if isinstance(data, pd.DataFrame) else None
+
+
+assert get_length(X_test) == get_length(
+    test_ids
+), f"Mismatch in length of test images and test IDs: X_test ({get_length(X_test)}) and test_ids ({get_length(test_ids)})"
+assert get_length(X) == get_length(
+    y
+), f"Mismatch in length of training images and labels: X ({get_length(X)}) and y ({get_length(y)})"
+
+assert get_length(X) != 0, f"Training data is empty."
+assert get_length(y) != 0, f"Training labels are empty."
+assert get_length(X_test) != 0, f"Test data is empty."
+
+assert get_width(X) == get_width(
+    X_test
+), "Mismatch in width of training and test data. Width means the number of features."
+
+if isinstance(X, pd.DataFrame) and isinstance(X_test, pd.DataFrame):
+    assert get_column_list(X) == get_column_list(X_test), "Mismatch in column names of training and test data."
+
+assert get_width(X) == get_width(
+    X_test
+), "Mismatch in width of training and test data. Width means the number of features."
 
 print("Data loader test passed successfully. Length of test images matches length of test IDs.")

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

@@ -373,6 +373,13 @@ data_loader_eval:
     {{test_code}}
     ```
 
+    {% if workflow_stdout is not none %}
+    Your feature engineering code is also part of the whole workflow, the user also tested the whole workflow and provided you the stdout.
+    The whole workflow code is:
+    {{workflow_code}}
+    Please consider both stdout and approve the code when both the feature engineering test and the whole workflow test pass.
+    {% endif %}
+    
     You'll be given the stdout of your testing scripts.
     Please respond with your feedback in the following JSON format and order
     ```json
@@ -384,6 +391,9 @@ data_loader_eval:
     }
     ```
   user: |-
-    ```
-    {{stdout}}
-    ```
+    Data loader test stdout:
+    {{stdout}}   
+    {% if workflow_stdout is not none %}
+    Whole workflow test stdout:
+    {{workflow_stdout}}
+    {% endif %}