Skip to content

magics

Magic to generate code cells for notebooks using OpenAI's API.

assist(line, cell) #

Generate code cells for notebooks using OpenAI's API.

Usage:

%%assist
# create a scatterplot from df

Will create a code cell below the current one looking something like this

# create a scatterplot from df

df.plot.scatter(x="col1", y="col2")

If you've run previous cells in the notebook, the generated code will be based on the history of the current notebook session.

That even means that, for example, running df.columns in a previous cell will cause generated code to use column names in the future!

Caveats:

  • Only the last 5 cell executions are provided as context.
  • The generated code is not guaranteed to be correct, idiomatic, efficient, readable, or useful.
  • The generated code is not guaranteed to be syntactically correct or even something to write home about.

There are several options you can use to control the behavior of the magic.

--fresh: Ignore the history of the current notebook session and generate code from scratch.

--in-place: Replace the current cell with the generated code.

--verbose: Show additional information in the cell output.

Example:

%%assist --fresh --in-place --verbose
# how can I query for pokemon via the PokéAPI?
Source code in genai/magics.py
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
@magic_arguments()
@argument("--fresh", action="store_true")
@argument(
    "--verbose",
    action="store_true",
    help="Show additional information in the cell output",
)
@argument(
    "--in-place",
    action="store_true",
    help="Replace the current cell with the generated code",
)
@cell_magic
def assist(line, cell):
    """Generate code cells for notebooks using OpenAI's API.

    Usage:

    ```python
    %%assist
    # create a scatterplot from df
    ```

    Will create a code cell below the current one looking something like this

    ```python
    # create a scatterplot from df

    df.plot.scatter(x="col1", y="col2")
    ```

    If you've run previous cells in the notebook, the generated code will be
    based on the history of the current notebook session.

    That even means that, for example, running `df.columns` in a previous cell
    will cause generated code to use column names in the future!

    Caveats:

    - Only the last 5 cell executions are provided as context.
    - The generated code is not guaranteed to be correct, idiomatic, efficient, readable, or useful.
    - The generated code is not guaranteed to be syntactically correct or even something to write home about.

    There are several options you can use to control the behavior of the magic.

    --fresh: Ignore the history of the current notebook session and generate code from scratch.

    --in-place: Replace the current cell with the generated code.

    --verbose: Show additional information in the cell output.

    Example:

    ```python
    %%assist --fresh --in-place --verbose
    # how can I query for pokemon via the PokéAPI?
    ```
    """
    progress = display(starting_message(), display_id=True)

    args = parse_argstring(assist, line)

    ip = get_ipython()
    cell_text = cell.strip()

    context = []
    if not args.fresh:
        context = get_historical_context(ip)

    if args.verbose:
        print("magic arguments:", line)
        print("submission:", cell)
        print("context:", context)

    # Pass streaming as False since we cannot replace after a `set_next_input`
    generated_text = generate_next_cell(context, cell_text, stream=False)

    progress.update(completion_made())

    preamble = ""

    if args.in_place:
        # Since we're running it in place, keep the context of what was sent in.
        # The preamble is a comment with the magic line and the original cell text all commented out
        processed_cell_text = "\n".join(f"# {line}" for line in cell_text.splitlines())
        preamble = f"""#%%assist {line}\n{processed_cell_text}\n"""

    new_cell = preamble

    for delta in generated_text:
        new_cell = new_cell + delta

    ip.set_next_input(
        new_cell,
        replace=args.in_place,
    )