feat: create Jupyter notebook pipeline file based on main.py file (#1134)

* First commit

* isort

* black

* tweak prompt

* fix for argparse

* fix typo

* add e2e

* Add test files, clean

* fix black settings

* revert

* fix trailing

* remove extra

* comment

* small fix, updated prompt

* Fix argparse

* small improvements

* fix for merge

* fix for merge

rdagent/components/coder/data_science/pipeline/__init__.py

@@ -83,7 +83,10 @@ class PipelineMultiProcessEvolvingStrategy(MultiProcessEvolvingStrategy):
             package_info=target_task.package_info,
             enable_model_dump=DS_RD_SETTING.enable_model_dump,
             enable_debug_mode=DS_RD_SETTING.sample_data_by_LLM,
-            spec=T("scenarios.data_science.share:component_spec.Pipeline").r(metric_name=self.scen.metric_name),
+            spec=T("scenarios.data_science.share:component_spec.Pipeline").r(
+                metric_name=self.scen.metric_name,
+                enable_notebook_conversion=DS_RD_SETTING.enable_notebook_conversion,
+            ),
         )
         user_prompt = T(".prompts:pipeline_coder.user").r(
             competition_info=competition_info,

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

@@ -16,6 +16,7 @@ from rdagent.components.coder.CoSTEER.knowledge_management import (
     CoSTEERQueriedKnowledgeV2,
 )
 from rdagent.components.coder.data_science.conf import get_clear_ws_cmd, get_ds_env
+from rdagent.components.coder.data_science.share.notebook import NotebookConverter
 from rdagent.components.coder.data_science.utils import remove_eda_part
 from rdagent.core.experiment import FBWorkspace, Task
 from rdagent.scenarios.data_science.test_eval import get_test_eval
@@ -70,6 +71,24 @@ class PipelineCoSTEEREvaluator(CoSTEEREvaluator):
                 env=env, entry=f"strace -e trace=file -f -o trace.log python -m coverage run main.py"
             )
 
+        nb_conversion_ret_code = 0
+        nb_conversion_check_text = ""
+        if DS_RD_SETTING.enable_notebook_conversion:
+            notebook_converter = NotebookConverter()
+            code = implementation.file_dict["main.py"]
+            error_msg = notebook_converter.validate_code_format(code)
+            if error_msg is not None:
+                nb_conversion_check_text = error_msg
+                nb_conversion_ret_code = 1
+            else:
+                notebook_converter.convert(
+                    task=target_task,
+                    code=code,
+                    stdout=result.stdout,
+                    outfile=implementation.workspace_path / "main.ipynb",
+                    use_debug_flag=DS_RD_SETTING.sample_data_by_LLM,
+                )
+
         sample_submission_check = True
         test_eval = get_test_eval()
         if (sample_submission_file_name := test_eval.get_sample_submission_name(self.scen.competition)) is not None:
@@ -173,7 +192,10 @@ class PipelineCoSTEEREvaluator(CoSTEEREvaluator):
             scenario=self.scen.get_scenario_all_desc(eda_output=eda_output),
             task_desc=target_task.get_task_information(),
             stdout=stdout.strip(),
-            spec=T("scenarios.data_science.share:component_spec.Pipeline").r(metric_name=self.scen.metric_name),
+            spec=T("scenarios.data_science.share:component_spec.Pipeline").r(
+                metric_name=self.scen.metric_name,
+                enable_notebook_conversion=DS_RD_SETTING.enable_notebook_conversion,
+            ),
             code=implementation.file_dict["main.py"],
         )
         wfb = build_cls_from_json_with_retry(
@@ -193,4 +215,7 @@ class PipelineCoSTEEREvaluator(CoSTEEREvaluator):
             wfb.return_checking += (
                 "\nSample submission file check failed. Code should not open the sample submission file."
             )
+        if nb_conversion_ret_code != 0 and wfb.final_decision is True:
+            wfb.final_decision = False
+            wfb.return_checking += "\n" + nb_conversion_check_text
         return wfb

rdagent/components/coder/data_science/share/notebook.py

@@ -0,0 +1,135 @@
+"""
+Handles conversion from a Python file to a Jupyter notebook.
+"""
+
+import argparse
+from typing import Optional
+
+import nbformat
+
+from rdagent.components.coder.data_science.share.util import (
+    extract_first_section_name_from_code,
+    extract_function_body,
+    split_code_and_output_into_sections,
+)
+from rdagent.core.experiment import Task
+from rdagent.log import rdagent_logger as logger
+from rdagent.oai.llm_utils import APIBackend
+from rdagent.utils.agent.ret import MarkdownAgentOut
+from rdagent.utils.agent.tpl import T
+
+
+class NotebookConverter:
+    """
+    Builder responsible for writing a Jupyter notebook for a workspace.
+    """
+
+    def validate_code_format(self, code: str) -> str | None:
+        """
+        Returns None if the code format is valid, otherwise returns an error message.
+        """
+        main_function_body = extract_function_body(code, "main")
+        if not main_function_body:
+            return "[Error] No main function found in the code. Please ensure that the main function is defined and contains the necessary print statements to divide sections."
+
+        found_section_name = extract_first_section_name_from_code(main_function_body)
+        if not found_section_name:
+            return "[Error] No sections found in the code. Expected to see 'print(\"Section: <section name>\")' as section dividers. Also make sure that they are actually run and not just comments."
+
+        return None
+
+    def convert(
+        self,
+        task: Optional[Task],
+        code: str,
+        stdout: str,
+        outfile: Optional[str] = None,
+        use_debug_flag: bool = False,
+    ) -> str:
+        """
+        Build a notebook based on the current progression.
+        """
+        # Handle argparse in the code to ensure it works in a notebook environment
+        should_handle_argparse = "argparse" in code
+        sections = split_code_and_output_into_sections(code=code, stdout=stdout)
+        notebook = nbformat.v4.new_notebook()
+
+        # Use LLM to generate an intro cell for the notebook
+        if task:
+            system_prompt = T(".prompts:notebookconverter.system").r()
+            user_prompt = T(".prompts:notebookconverter.user").r(
+                plan=task.get_task_information(),
+                code=code,
+            )
+            resp = APIBackend().build_messages_and_create_chat_completion(
+                user_prompt=user_prompt, system_prompt=system_prompt
+            )
+            intro_content = MarkdownAgentOut.extract_output(resp)
+            notebook.cells.append(nbformat.v4.new_markdown_cell(intro_content))
+
+        if should_handle_argparse:
+            # Remove extra `import sys` since it will be added for argparse handling
+            if "import sys\n" in sections[0]["code"]:
+                sections[0]["code"] = sections[0]["code"].replace("import sys\n", "")
+
+            # Add sys.argv modification for argparse handling
+            sections[0]["code"] = (
+                "\n".join(
+                    [
+                        "import sys",
+                        "# hack to allow argparse to work in notebook",
+                        ('sys.argv = ["main.py", "--debug"]' if use_debug_flag else 'sys.argv = ["main.py"]'),
+                    ]
+                )
+                + "\n\n"
+                + sections[0]["code"].lstrip()
+            )
+
+        for section in sections:
+            # Create a markdown cell for the section name and comments
+            markdown_content = ""
+            if section["name"]:
+                markdown_content += f"## {section['name']}\n"
+            if section["comments"]:
+                markdown_content += f"{section['comments']}\n"
+            if markdown_content:
+                notebook.cells.append(nbformat.v4.new_markdown_cell(markdown_content))
+
+            # Create a code cell for the section code and output
+            if section["code"]:
+                cell = nbformat.v4.new_code_cell(section["code"])
+                if section["output"]:
+                    # For simplicity, treat all output as coming from stdout
+                    # TODO: support Jupyter kernel execution and handle outputs appropriately here
+                    cell.outputs = [nbformat.v4.new_output("stream", name="stdout", text=section["output"])]
+                notebook.cells.append(cell)
+
+        # Save the notebook or return it as a string
+        if outfile:
+            with open((outfile), "w", encoding="utf-8") as f:
+                nbformat.write(notebook, f)
+                logger.info(f"Notebook written to {outfile}")
+
+        return nbformat.writes(notebook)
+
+
+if __name__ == "__main__":
+    converter = NotebookConverter()
+    parser = argparse.ArgumentParser(description="Convert Python code to Jupyter notebook.")
+    parser.add_argument("inputfile", type=str, help="Path to the input Python file.")
+    parser.add_argument("outfile", type=str, help="Path to the output Notebook file.")
+    parser.add_argument(
+        "--stdout",
+        type=str,
+        default="",
+        help="Standard output from the code execution.",
+    )
+    parser.add_argument("--debug", action="store_true", help="Use debug flag to modify sys.argv.")
+    args = parser.parse_args()
+    converter.convert(
+        task=None,
+        code=open(args.inputfile, "r").read(),
+        stdout=args.stdout,
+        outfile=args.outfile,
+        use_debug_flag=False,
+    )

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

@@ -89,3 +89,25 @@ docdev:
     ```
     {% endfor %}
 
+notebookconverter:
+  system: |-
+    {% include "scenarios.data_science.share:scen.role" %} Your task is to provide a summary for a data science solution.
+
+    You will be given:
+    - The original implementation plan for the script.
+    - A Python script that contains code and output.
+
+    Your task is to generate markdown content that includes a title and a short paragraph summarizing the technique in model training, the type of model produced and any other noteworthy details in the solution.
+
+    The return content should be like the format below(Please note that "````" is used to avoid confliction of "```" in markdown file)
+    ````markdown
+    # <The title of the notebook>
+    <the content of markdown file>
+    ````
+
+  user: |-
+    --------------- The implementation plan ---------------
+    {{plan}}
+
+    --------------- The Python script content ---------------
+    {{code}}

rdagent/components/coder/data_science/share/util.py

@@ -0,0 +1,365 @@
+import ast
+import io
+import re
+import tokenize
+from itertools import zip_longest
+from typing import List, Optional, Set, Tuple, TypedDict
+
+
+class CodeSection(TypedDict):
+    """
+    Represents a section of the original Python source code, to be converted to a notebook cell.
+    """
+
+    name: Optional[str]
+    code: Optional[str]
+    comments: Optional[str]
+    output: Optional[str]
+
+
+def extract_function_body(source_code: str, function_name: str) -> Optional[str]:
+    """
+    Extracts the body of a function from the source code.
+    Returns None if the function is not found.
+
+    Assumption: The function is multiline and defined at the top level.
+    """
+    tree = ast.parse(source_code)
+    for node in ast.walk(tree):
+        if isinstance(node, ast.FunctionDef) and node.name == function_name:
+            lines = source_code.splitlines()
+            start = node.body[0].lineno
+            end = node.body[-1].end_lineno
+            body_lines = lines[start - 1 : end]
+            indent_level = len(body_lines[0]) - len(body_lines[0].lstrip())
+            return "\n".join(line[indent_level:] for line in body_lines)
+    return None
+
+
+def split_sections(
+    text: str, section_header_regex: str, known_sections: Optional[list[str]] = None
+) -> tuple[Optional[str], list[str], list[str]]:
+    """
+    Split text into sections based on the section headers.
+    """
+    sections = []
+    section_names = []
+    current_section = []
+    next_section_name_index = 0
+    for line in text.splitlines():
+        match = re.match(section_header_regex, line)
+        extracted_section_name = match.group(1).strip() if match else None
+        if extracted_section_name and (
+            not known_sections
+            or (
+                next_section_name_index < len(known_sections)
+                and extracted_section_name == known_sections[next_section_name_index]
+            )
+        ):
+            if current_section:
+                sections.append("\n".join(current_section))
+                current_section = []
+            current_section.append(line)
+            section_names.append(extracted_section_name)
+            next_section_name_index += 1
+        else:
+            current_section.append(line)
+    if current_section:
+        sections.append("\n".join(current_section))
+
+    # If the first section does not match the header regex, treat it as a header section.
+    header_section = None
+    if sections and not re.search(section_header_regex, sections[0]):
+        header_section = sections[0]
+        sections = sections[1:]
+
+    return header_section, sections, section_names
+
+
+def split_code_sections(source_code: str) -> tuple[Optional[str], list[str]]:
+    """
+    Split code into sections based on the section headers.
+    """
+    return split_sections(source_code, r'^print\(["\']Section: (.+)["\']\)')
+
+
+def split_output_sections(stdout: str, known_sections: list[str]) -> tuple[Optional[str], list[str]]:
+    """
+    Split output into sections based on the section headers.
+    """
+    header_section, sections, _ = split_sections(stdout, r"^Section: (.+)", known_sections=known_sections)
+    return header_section, sections
+
+
+def extract_comment_under_first_print(source_code) -> tuple[Optional[str], str]:
+    """
+    Extract comments from the source code after the first print statement.
+    """
+    lines = source_code.splitlines()
+    lines_to_remove = set()
+    all_comments = []
+
+    parsed = ast.parse(source_code)
+    # Find the first print statement only
+    first_print_lineno = None
+    for node in ast.walk(parsed):
+        if isinstance(node, ast.Expr) and isinstance(node.value, ast.Call):
+            if getattr(node.value.func, "id", None) == "print":
+                first_print_lineno = node.lineno
+                break
+
+    if first_print_lineno is None:
+        # No print statement found, return empty comments and original code
+        return None, source_code
+
+    for i in range(first_print_lineno, len(lines)):
+        stripped = lines[i].strip()
+        if stripped.startswith("#"):
+            comment_text = stripped.lstrip("# ").strip()
+            all_comments.append(comment_text)
+            lines_to_remove.add(i)
+        elif stripped == "":
+            continue
+        elif i > first_print_lineno:
+            break  # stop after hitting actual code line
+
+    cleaned_lines = [line for idx, line in enumerate(lines) if idx not in lines_to_remove]
+    cleaned_code = "\n".join(cleaned_lines)
+    comments_str = "\n".join(all_comments) if all_comments else None
+
+    return comments_str, cleaned_code
+
+
+def extract_first_section_name_from_code(source_code):
+    """
+    Extract the first section name from the source code.
+    """
+    parsed = ast.parse(source_code)
+    for node in ast.walk(parsed):
+        if isinstance(node, ast.Expr) and isinstance(node.value, ast.Call):
+            call = node.value
+            if getattr(call.func, "id", None) == "print" and call.args:
+                arg0 = call.args[0]
+                if isinstance(arg0, ast.Constant) and isinstance(arg0.value, str):
+                    # Match "Section: ..." pattern
+                    m = re.match(r"Section:\s*(.+)", arg0.value)
+                    if m:
+                        return m.group(1).strip()
+    return None
+
+
+def extract_first_section_name_from_output(stdout: str) -> Optional[str]:
+    """
+    Extract the first section name from the output string.
+    """
+    match = re.search(r"Section:\s*(.+)", stdout)
+    if match:
+        return match.group(1).strip()
+    return None
+
+
+def is_function_called(source_code: str, func_name: str) -> bool:
+    """
+    Returns True if the function named `func_name` is called in `source_code`.
+    """
+    tree = ast.parse(source_code)
+    for node in ast.walk(tree):
+        if isinstance(node, ast.Call):
+            # For simple function calls like func()
+            if isinstance(node.func, ast.Name) and node.func.id == func_name:
+                return True
+
+            # For calls like module.func()
+            elif isinstance(node.func, ast.Attribute) and node.func.attr == func_name:
+                return True
+    return False
+
+
+def remove_function(source_code: str, function_name: str) -> str:
+    """
+    Remove a function definition from the source code.
+    """
+    tree = ast.parse(source_code)
+    lines = source_code.splitlines()
+
+    for node in tree.body:
+        if isinstance(node, ast.FunctionDef) and node.name == function_name:
+            start_lineno = node.lineno - 1
+            end_lineno = node.end_lineno
+            return "\n".join(lines[:start_lineno] + lines[end_lineno:])
+
+    return source_code
+
+
+def remove_main_block(source_code: str) -> str:
+    """
+    Remove the if __name__ == "__main__": block from the source code.
+    """
+    tree = ast.parse(source_code)
+    lines = source_code.splitlines()
+
+    # Find the main block and note its line numbers
+    for node in tree.body:
+        if isinstance(node, ast.If):
+            test = node.test
+            if (
+                isinstance(test, ast.Compare)
+                and isinstance(test.left, ast.Name)
+                and test.left.id == "__name__"
+                and len(test.ops) == 1
+                and isinstance(test.ops[0], ast.Eq)
+                and len(test.comparators) == 1
+                and isinstance(test.comparators[0], ast.Constant)
+                and test.comparators[0].value == "__main__"
+            ):
+
+                # Remove lines corresponding to this block
+                start_lineno = node.lineno - 1
+                end_lineno = node.end_lineno
+                return "\n".join(lines[:start_lineno] + lines[end_lineno:])
+
+    return source_code
+
+
+def extract_top_level_functions_with_decorators_and_comments(
+    code: str,
+) -> List[Tuple[str, str]]:
+    """
+    Returns list of (function_name, source_segment) for top-level functions (excluding "main"),
+    including decorators and contiguous preceding comments.
+    """
+    # Parse AST to get function nodes
+    tree = ast.parse(code)
+    lines = code.splitlines(keepends=True)
+
+    # Precompute which line numbers have comment tokens
+    comment_lines: Set[int] = set()
+    lines = code.splitlines(keepends=True)  # preserve exact line content for prefix checks
+
+    tokgen = tokenize.generate_tokens(io.StringIO(code).readline)  # yields (type, string, start, end, line)
+    for tok_type, _, (srow, scol), _, _ in tokgen:
+        if tok_type == tokenize.COMMENT:
+            # everything before the comment on that line must be whitespace
+            prefix = lines[srow - 1][:scol]
+            if prefix.strip() == "":
+                comment_lines.add(srow)
+
+    functions = []
+
+    for node in tree.body:  # only top-level
+        if not isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
+            continue
+        if node.name == "main":
+            continue
+
+        # Determine the starting line: earliest decorator if present, else the def/async line
+        if node.decorator_list:
+            start_lineno = min(d.lineno for d in node.decorator_list)
+        else:
+            start_lineno = node.lineno
+
+        # Extend upward to include contiguous comment lines (no intervening non-blank/non-comment)
+        span_start = start_lineno
+        curr = span_start - 1  # check line above; lines are 1-based
+        while curr > 0:
+            line_text = lines[curr - 1]
+            if curr in comment_lines:
+                span_start = curr
+                curr -= 1
+                continue
+            if line_text.strip() == "":
+                # blank line: include it and keep scanning upward
+                span_start = curr
+                curr -= 1
+                continue
+            break  # encountered code or something else; stop
+
+        # Determine end line of the function definition including its body
+        # Prefer end_lineno if available (Python 3.8+)
+        if hasattr(node, "end_lineno") and node.end_lineno is not None:
+            span_end = node.end_lineno
+        else:
+            # Fallback: get last lineno from the deepest child in body
+            def _max_lineno(n):
+                max_ln = getattr(n, "lineno", 0)
+                for child in ast.iter_child_nodes(n):
+                    ln = _max_lineno(child)
+                    if ln > max_ln:
+                        max_ln = ln
+                return max_ln
+
+            span_end = _max_lineno(node)
+
+        # Slice the original source lines
+        segment = "".join(lines[span_start - 1 : span_end])
+        functions.append((node.name, segment))
+
+    return functions
+
+
+def split_code_and_output_into_sections(code: str, stdout: str) -> list[CodeSection]:
+    """
+    Converts a Python script and its output into a list of CodeSections.
+    Pre-condition: The code in the main() function contains print statements that indicate section names, e.g., `print("Section: <section name>")`.
+    """
+    # This will hold all top-level code and by default all function definitions.
+    # Functions will later be moved to more relevant sections if needed.
+    # The first step is to remove both the if __name__ == "__main__": block and the main function
+    top_level_code = remove_main_block(remove_function(code, "main"))
+
+    main_function_body = extract_function_body(code, "main")
+    functions = extract_top_level_functions_with_decorators_and_comments(top_level_code)
+
+    # Split the main function body into sections based on print("Section: <section name>") code
+    main_fn_top_level_section, main_fn_sections, known_section_names = (
+        split_code_sections(main_function_body) if main_function_body else (None, [], [])
+    )
+
+    # Split the output into sections based on "Section: " headers
+    output_top_level_section, output_sections = split_output_sections(stdout, known_section_names)
+
+    # Merge code and outputs into code sections
+    result_sections: list[CodeSection] = []
+    for output_section, code_section in zip_longest(output_sections, main_fn_sections):
+        name = None
+        if code_section is not None:
+            # If code section is available, extract the section name from it
+            name = extract_first_section_name_from_code(code_section)
+        elif output_section:
+            # If only output section is available, extract the section name from it
+            name = extract_first_section_name_from_output(output_section)
+        comments, cleaned_code = (
+            extract_comment_under_first_print(code_section) if code_section is not None else (None, None)
+        )
+        # Strip whitespaces for the cell
+        if cleaned_code is not None:
+            cleaned_code = cleaned_code.strip()
+        result_sections.append(CodeSection(name=name, code=cleaned_code, comments=comments, output=output_section))
+
+    # Small optimization: move function definitions to the sections where they are first called
+    # TODO: this doesn't handle nested function references, e.g., fn A calls fn B which calls fn C
+    # currently will not move C to the section where A is called
+    for name, segment in functions:
+        for section in result_sections:
+            if section["code"] and is_function_called(section["code"], name):
+                section["code"] = segment.strip() + "\n\n" + section["code"].lstrip()
+                top_level_code = top_level_code.replace(segment, "")
+                break
+
+    # Inject the top-level code at the beginning of the sections
+    top_level_code = (
+        top_level_code.rstrip() + "\n\n" + main_fn_top_level_section.lstrip()
+        if main_fn_top_level_section
+        else top_level_code
+    )
+    result_sections.insert(
+        0,
+        CodeSection(
+            name=None,
+            code=top_level_code,
+            comments=None,
+            output=output_top_level_section,
+        ),
+    )
+
+    return result_sections