Merge pull request #237 from MLSysOps/feat/improve-kaggle

[MRG] Fixed bugs of Kaggle integration

mle/agents/advisor.py

@@ -55,7 +55,7 @@ def __init__(self, model, console=None):
          the model (or method), and the evaluation metrics, etc. You should always follow the user's requirements.
         2. You should briefly analyze the user's dataset, and give a summary of the dataset, the dataset input can be
          a public dataset name or a path to a local CSV file. You can use the function `preview_csv_data` to preview
-         the CSV file or not if the dataset is a public dataset.
+         the CSV files or not if the dataset is a public dataset.
         3. And then you should always use the function `search_arxiv` or `search_papers_with_code` to search the
          state-of-the-art machine learning tasks/models/algorithms that can be used to solve the user's requirements,
           and stay up-to-date with the latest.
@@ -102,11 +102,12 @@ def __init__(self, model, console=None):
         self.sys_prompt += self.json_mode_prompt
         self.chat_history.append({"role": 'system', "content": self.sys_prompt})
 
-    def suggest(self, requirement):
+    def suggest(self, requirement, process=True):
         """
         Handle the query from the model query response.
         Args:
             requirement: the user requirement.
+            process: whether to process the report.
         """
         with self.console.status("MLE Advisor is thinking the best strategy to help you..."):
             self.chat_history.append({"role": "user", "content": requirement})
@@ -123,6 +124,9 @@ def suggest(self, requirement):
             except json.JSONDecodeError as e:
                 suggestions = clean_json_string(text)
 
+        if not process:
+            return suggestions
+
         return process_report(requirement, suggestions)
 
     def interact(self, requirement):
@@ -158,7 +162,6 @@ def interact(self, requirement):
 
         return self.report
 
-
     def clarify_dataset(self, dataset: str):
         """
         Clarify the dataset blurity by suggest some datasets.

mle/agents/coder.py

@@ -24,14 +24,17 @@ def process_summary(summary_dict: dict):
 
 class CodeAgent:
 
-    def __init__(self, model, working_dir='.', console=None):
+    def __init__(self, model, working_dir='.', console=None, single_file=False):
         """
         CodeAgent: the agent to solve the given coding problems, by planning coding tasks, searching websites,
         and generating code snippets. It does not execute the code, only make use of built-in functions to provides
          the code snippets to solve the problem.
 
         Args:
             model: the model to use.
+            working_dir: the working directory.
+            console: the console to use.
+            single_file: whether the agent is working on a single file or not.
         """
         config_data = get_config()
         self.code_summary = None
@@ -49,19 +52,37 @@ def __init__(self, model, working_dir='.', console=None):
 
         Your can leverage your capabilities by using the specific functions listed below:
 
-        1. Creating project structures based on the user requirement using function `create_directory`.
-        2. Writing clean, efficient, and well-documented code using function `create_file` and `write_file`.
-        3. Exam the project to re-use the existing code snippets as much as possible, you may need to use
+        - Creating project structures based on the user requirement using function `create_directory`.
+        - Writing clean, efficient, and well-documented code using function `create_file` and `write_file`.
+        - Exam the project to re-use the existing code snippets as much as possible, you may need to use
          functions like `list_files`, `read_file` and `write_file`.
-        4. Writing the code into the file when creating new files, do not create empty files.
-        5. Use function `preview_csv_data` to preview the CSV data if the task include CSV data processing.
-        6. Decide whether the task requires execution and debugging before moving to the next or not.
-        7. Generate the commands to run and test the current task, and the dependencies list for this task.
-        8. You only write Python scripts, don't write Jupiter notebooks which require interactive execution.
+        - Writing the code into the file when creating new files, do not create empty files.
+        - Use function `preview_csv_data` to preview the CSV data if the task include CSV data processing.
+        - Decide whether the task requires execution and debugging before moving to the next or not.
+        - Generate the commands to run and test the current task, and the dependencies list for this task.
+        - You only write Python scripts, don't write Jupiter notebooks which require interactive execution.
         """
+
+        if single_file:
+            self.sys_prompt = f"""
+            You are an expert programmer working on an Machine Learning task using Python.
+            You are currently working on: {self.working_dir}.
+            
+            Your can leverage your capabilities by using the specific functions listed below:
+            
+            - You should create a single script first, with the complete code inside. You can have multiple functions and classes.
+            - Writing clean, efficient, and well-documented code to a script using functions `create_file`.
+            - Use function `preview_csv_data` to preview the CSV data if the task include CSV dataset or examples.
+            - Generate the commands to run and test the current script, and the dependencies list required for this script.
+            - You only write Python scripts, don't write Jupiter notebooks which require interactive execution.
+            - Make sure the code has met the task description, and the suggested methods.
+            - Make sure the output format and the output file path is correct.
+            """
+
         self.search_prompt = """
-        9. Performing web searches use function `web_search` to get up-to-date information or additional context.
+        - Performing web searches use function `web_search` to get up-to-date information or additional context.
         """
+
         self.json_mode_prompt = """
 
         The output format should be in JSON format, include:
@@ -86,6 +107,27 @@ def __init__(self, model, working_dir='.', console=None):
         }
         
         """
+
+        if single_file:
+            self.json_mode_prompt = """
+        
+            The output format should be in JSON format, include:
+            
+            1. The dependency list that the project needs to run.
+            2. And the command and the parameters to run and test the script.
+            
+            Example JSON output:
+            
+            {
+               "dependency":[
+                  "torch",
+                  "scikit-learn"
+               ],
+               "command":"python /path/to/your/project.py",
+            }
+            
+            """
+
         self.functions = [
             schema_read_file,
             schema_create_file,
@@ -122,7 +164,7 @@ def code(self, task_dict: dict):
         """
         task_prompt = textwrap.dedent(f"""
             You are required to complete task: {task_dict.get('task')}.\n
-            You should: {task_dict.get('description')}
+            Task description: {task_dict.get('description')}
             """)
 
         with self.console.status(f"Coder is working on the task: {task_dict.get('task')}..."):

mle/agents/debugger.py

@@ -27,14 +27,15 @@ def process_debug_report(debug_report):
 
 class DebugAgent:
 
-    def __init__(self, model, console=None):
+    def __init__(self, model, console=None, analyze_only=False):
         """
         DebugAgent: the agent to run the generated the code and then debug it. The return of the
         agent is an instruction to the user to modify the code based on the logs and web search.
 
         Args:
             model: the model to use.
             console: the console to use.
+            analyze_only: if only analyze the code without execution
         """
         config_data = get_config()
         self.console = console
@@ -47,21 +48,21 @@ def __init__(self, model, console=None):
 
         Your can leverage your capabilities by using the specific functions listed below:
 
-        1. Install the code dependencies using function `execute_command` based on the Developer's dependencies list.
-        2. Execute the code using function `execute_command` to test the code based on the Developer's instructions.
-        3. If the program returns errors, you need to debug the code based on the logs, you may need to first read the
+        - Install the code dependencies using function `execute_command` based on the Developer's dependencies list.
+        - Execute the code using function `execute_command` to test the code based on the Developer's instructions.
+        - If the program returns errors, you need to debug the code based on the logs, you may need to first read the
          structure of the project using function `list_files`.
-        4. Then you may need to call `read_file` function to read the content of the code files, locate the error line
+        - Then you may need to call `read_file` function to read the content of the code files, locate the error line
          and the reasons.
-        5. You don't need to care about the best practices and code styles, you only care about the errors in the code.
-        6. If the developer's input includes the error message, you don't need to execute the code, you can directly
-         analyze the error messages and give the reason and suggestions.
+        - You don't need to care about the best practices and code styles, you only care about the errors in the code.
 
         """
+
         self.search_prompt = """
-        7. You need to debug the code based on the error logs, you may need to call `web_search` function to search for
+        - You need to debug the code based on the error logs, you may need to call `web_search` function to search for
          the solutions or reasons for the errors if needed.
         """
+
         self.json_mode_prompt = """
         
         Example JSON output if a program runs without errors:
@@ -84,34 +85,65 @@ def __init__(self, model, console=None):
             "suggestion":"Failed to find the target file. Please check the file path."
            ]
         }
-        
-        Example JSON output if the error message is provided (no running status required):
-        {
-           "status":"",
-           "changes":[
-               {
-                     "file":"xxx.py",
-                     "line":15,
-                     "issue":"xxx",
-                     "suggestion":"xxx"
-               }
-           ],
-           "suggestion":"Incorrect file imports. Please replace the import statement with 'from xxx import xxx'."
-        }
         """
+
         self.functions = [
             schema_read_file,
             schema_list_files,
             schema_execute_command
         ]
 
+        if analyze_only:
+            self.sys_prompt = """
+            You are a program error debugger working on a Python project. You target is to
+             analyze the running logs and the source code to locate the errors. And give the
+             suggestions to the developer to fix the errors.
+
+            Your can leverage your capabilities by using the specific functions listed below:
+
+            - Read and understand the running logs and the error messages.
+            - Read the content of the source code files using function `read_file`, to locate the error line.
+            - Give the suggestions to the developer to fix the errors.
+            - Install missing dependencies using function `execute_command` based on the error logs.
+            - If there is no error based on the exit code, you don't need to do anything but return the success status.
+            """
+
         if config_data.get('search_key'):
             self.functions.append(schema_web_search)
             self.sys_prompt += self.search_prompt
 
         self.sys_prompt += self.json_mode_prompt
         self.chat_history.append({"role": 'system', "content": self.sys_prompt})
 
+    def analyze_with_log(self, commands, logs):
+        """
+        Analyze the logs.
+        :param commands: the commands to execute.
+        :param logs: the logs to analyze.
+        :return:
+        """
+        analyze_prompt = f"""
+        The command to execute the code: {commands} \n
+        The logs to analyze: {logs} \n
+        """
+
+        self.chat_history.append({"role": "user", "content": analyze_prompt})
+        try:
+            text = self.model.query(
+                self.chat_history,
+                function_call='auto',
+                functions=self.functions,
+                response_format={"type": "json_object"}
+            )
+        except Exception as e:
+            print(f"Error occurred while querying the model: {e}")
+            return {}
+
+        self.chat_history.append({"role": "assistant", "content": text})
+        report_dict = json.loads(text)
+        print_in_box(process_debug_report(report_dict), self.console, title="MLE Debugger", color="yellow")
+        return report_dict
+
     def analyze(self, code_report):
         """
         Handle the query from the model query response.
@@ -145,4 +177,4 @@ def analyze(self, code_report):
         self.chat_history.append({"role": "assistant", "content": text})
         report_dict = json.loads(text)
         print_in_box(process_debug_report(report_dict), self.console, title="MLE Debugger", color="yellow")
-        return report_dict
+        return report_dict

mle/agents/summarizer.py

@@ -134,24 +134,42 @@ def summarize(self):
 
         return summary
 
-    def kaggle_request_summarize(self, kaggle_overview):
+    def kaggle_request_summarize(
+            self,
+            kaggle_overview: str,
+            exp_submission: str = None,
+            submission_file: str = "./submission.csv"):
         """
         Summarize the kaggle requests.
         :params: kaggle_overview: the overview json of kaggle competition
+        :params: exp_submission: the expected sample submission file.
+        :params: submission_file: the submission file path.
         """
         system_prompt = """
         You are a seasoned data science expert in Kaggle competitions. Your task is to summarize the
         requirements of a specific Kaggle competition in a clear and concise manner. Please ensure that
         your summary includes the following aspects:
 
-        1. **Overview**: Describe the competition's objective and significance.
-        2. **Data**: Detail the datasets, including file types, structure, and key features.
-        3. **Evaluation**: Explain the judging metric and its calculation.
-        4. **Submission**: Outline the format and requirements for submissions.
-        5. **Rules**: Highlight important rules, including data usage, team composition, and resources.
+        1. Overview: Describe the competition's objective and significance.
+        2. Data: Detail the datasets, including file types, structure, and key features.
+        3. Evaluation: Explain the judging metric and its calculation.
+        4. Submission: Outline the format and requirements for submissions, you can includes some submission samples using function `preview_csv_data`.
+        5. Rules: Highlight important rules, including data usage, team composition, and resources.
         """
+
+        if exp_submission:
+            kaggle_overview += f"""
+            \n\n EXAMPLE SUBMISSION FILE: {exp_submission}\n
+            SUBMISSION FILE LOCATION: {submission_file}
+            """
+
         chat_history = [
             {"role": "system", "content": system_prompt},
-            {"role": "user", "content": str(kaggle_overview)}
+            {"role": "user", "content": kaggle_overview}
         ]
-        return self.model.query(chat_history)
+
+        return self.model.query(
+            chat_history,
+            function_call='auto',
+            functions=[schema_preview_csv_data]
+        )