Adds code format option (#57)

* Adds code format

* Adds code to docs

docs/source/users/index.md

@@ -128,6 +128,7 @@ By default, Jupyter AI assumes that a model will output markdown, so the output
 an `%%ai` command will be formatted as markdown by default. You can override this
 using the `-f` or `--format` argument to your magic command. Valid formats include:
 
+- `code`
 - `markdown`
 - `math`
 - `html`
@@ -149,6 +150,22 @@ will look like properly typeset equations.
 Generate the 2D heat equation in LaTeX surrounded by `$$`. Do not include an explanation.
 ```
 
+This prompt will produce output as a code cell below the input cell. 
+
+:::{warning}
+:name: run-code
+**Please review any code that a generative AI model produces before you run it
+or distribute it.**
+The code that you get in response to a prompt may have negative side effects and may
+include calls to nonexistent (hallucinated) helper functions.
+:::
+
+```
+%%ai chatgpt -f code
+A function that computes the lowest common multiples of two integers, and
+a function that runs 5 test cases of the lowest common multiple function
+```
+
 ### Interpolating IPython in prompts
 
 Using curly brace syntax, you can include variables and other IPython expressions in your

examples/code.ipynb

@@ -0,0 +1,130 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "c7c709c4-0dbe-4f82-8ee6-f4cf2fd55b68",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "%reload_ext jupyter_ai"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "f69abe97-654c-4e4a-8d18-575deca9c935",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "%%ai chatgpt --format code\n",
+    "A program that asks me for my name and then greets me by my name, in Polish"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "a832d3be-7a0b-49bc-aaf6-da8a32a89285",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdin",
+     "output_type": "stream",
+     "text": [
+      "Jak masz na imię?  foo\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Cześć foo!\n"
+     ]
+    }
+   ],
+   "source": [
+    "name = input(\"Jak masz na imię? \")\n",
+    "print(\"Cześć \" + name + \"!\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "7e187e37-9abd-4ed2-9959-fb3acdb3acfe",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%%ai chatgpt --format code\n",
+    "A function that computes the lowest common multiples of two integers, and a function that runs 5 test cases of the lowest common multiple function"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "555662c1-f576-4786-8dda-51ebe48c1d75",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "def lcm(x, y):\n",
+    "    if x > y:\n",
+    "        greater = x\n",
+    "    else:\n",
+    "        greater = y\n",
+    "\n",
+    "    while True:\n",
+    "        if (greater % x == 0) and (greater % y == 0):\n",
+    "            lcm = greater\n",
+    "            break\n",
+    "        greater += 1\n",
+    "\n",
+    "    return lcm\n",
+    "\n",
+    "def test_lcm():\n",
+    "    assert lcm(3, 5) == 15\n",
+    "    assert lcm(7, 9) == 63\n",
+    "    assert lcm(18, 24) == 72\n",
+    "    assert lcm(10, 15) == 30\n",
+    "    assert lcm(12, 16) == 48\n",
+    "\n",
+    "test_lcm()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "9978c08c-c0da-49dd-9e01-3b9d8c9e1a5e",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.9"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

packages/jupyter-ai-magics/jupyter_ai_magics/magics.py

@@ -20,6 +20,7 @@
 }
 
 DISPLAYS_BY_FORMAT = {
+    "code": None,
     "html": HTML,
     "markdown": Markdown,
     "math": Math,
@@ -31,6 +32,7 @@
 MARKDOWN_PROMPT_TEMPLATE = '{prompt}\n\nProduce output in markdown format only.'
 
 PROMPT_TEMPLATES_BY_FORMAT = {
+    "code": '{prompt}\n\nProduce output as source code only, with no text before or after it.',
     "html": '{prompt}\n\nProduce output in HTML format only, with no markup before or afterward.',
     "markdown": MARKDOWN_PROMPT_TEMPLATE,
     "md": MARKDOWN_PROMPT_TEMPLATE,
@@ -116,7 +118,7 @@ def _get_provider(self, provider_id: Optional[str]) -> BaseProvider:
                 optionally prefixed with the ID of the model provider, delimited
                 by a colon.""")
     @argument('-f', '--format',
-                choices=["markdown", "html", "json", "math", "md", "raw"],
+                choices=["code", "markdown", "html", "json", "math", "md", "raw"],
                 nargs="?",
                 default="markdown",
                 help="""IPython display to use when rendering output. [default="markdown"]""")
@@ -184,6 +186,17 @@ def ai(self, line, cell=None):
 
         # build output display
         DisplayClass = DISPLAYS_BY_FORMAT[args.format]
+
+        # if the user wants code, add another cell with the output.
+        if args.format == 'code':
+            new_cell_payload = dict(
+                source='set_next_input',
+                text=output,
+                replace=False,
+            )
+            ip.payload_manager.write_payload(new_cell_payload)
+            return None; # No output from the AI cell
+
         if DisplayClass is None:
             return output
         if args.format == 'json':