{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Jupyter Code Executor\n",
    "\n",
    "AutoGen is able to execute code in a stateful Jupyter kernel, this is in contrast to the command line code executor where each code block is executed in a new process. This means that you can define variables in one code block and use them in another. One of the interesting properties of this is that when an error is encountered, only the failing code needs to be re-executed, and not the entire script.\n",
    "\n",
    "To use the [`JupyterCodeExecutor`](/docs/reference/coding/jupyter/jupyter_code_executor#jupytercodeexecutor) you need a Jupyter server running. This can be in Docker, local, or even a remote server. Then, when constructing the [`JupyterCodeExecutor`](/docs/reference/coding/jupyter/jupyter_code_executor#jupytercodeexecutor) you pass it the server it should connect to.\n",
    "\n",
    "## Dependencies\n",
    "\n",
    "In order to use Jupyter based code execution some extra dependencies are required. These can be installed with the extra `jupyter-executor`:\n",
    "\n",
    "```bash\n",
    "pip install 'autogen-agentchat[jupyter-executor]~=0.2'\n",
    "```\n",
    "\n",
    "## Jupyter Server\n",
    "\n",
    "### Docker\n",
    "\n",
    "To run a Docker based Jupyter server, the [`DockerJupyterServer`](/docs/reference/coding/jupyter/docker_jupyter_server#dockerjupyterserver) can be used."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "exit_code=0 output='Hello, World!\\n' output_files=[]\n"
     ]
    }
   ],
   "source": [
    "from autogen.coding import CodeBlock\n",
    "from autogen.coding.jupyter import DockerJupyterServer, JupyterCodeExecutor\n",
    "\n",
    "with DockerJupyterServer() as server:\n",
    "    executor = JupyterCodeExecutor(server)\n",
    "    print(\n",
    "        executor.execute_code_blocks(\n",
    "            code_blocks=[\n",
    "                CodeBlock(language=\"python\", code=\"print('Hello, World!')\"),\n",
    "            ]\n",
    "        )\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "By default the [`DockerJupyterServer`](/docs/reference/coding/jupyter/docker_jupyter_server#dockerjupyterserver) will build and use a bundled Dockerfile, which can be seen below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "FROM quay.io/jupyter/docker-stacks-foundation\n",
      "\n",
      "SHELL [\"/bin/bash\", \"-o\", \"pipefail\", \"-c\"]\n",
      "\n",
      "USER ${NB_UID}\n",
      "RUN mamba install --yes jupyter_kernel_gateway ipykernel &&     mamba clean --all -f -y &&     fix-permissions \"${CONDA_DIR}\" &&     fix-permissions \"/home/${NB_USER}\"\n",
      "\n",
      "ENV TOKEN=\"UNSET\"\n",
      "CMD python -m jupyter kernelgateway --KernelGatewayApp.ip=0.0.0.0     --KernelGatewayApp.port=8888     --KernelGatewayApp.auth_token=\"${TOKEN}\"     --JupyterApp.answer_yes=true     --JupyterWebsocketPersonality.list_kernels=true\n",
      "\n",
      "EXPOSE 8888\n",
      "\n",
      "WORKDIR \"${HOME}\"\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(DockerJupyterServer.DEFAULT_DOCKERFILE)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Custom Docker Image\n",
    "\n",
    "A custom image can be used by passing the `custom_image_name` parameter to the [`DockerJupyterServer`](/docs/reference/coding/jupyter/docker_jupyter_server#dockerjupyterserver) constructor. There are some requirements of the image for it to work correctly:\n",
    "\n",
    "- The image must have [Jupyer Kernel Gateway](https://jupyter-kernel-gateway.readthedocs.io/en/latest/) installed and running on port 8888 for the [`JupyterCodeExecutor`](/docs/reference/coding/jupyter/jupyter_code_executor) to be able to connect to it.\n",
    "- Respect the `TOKEN` environment variable, which is used to authenticate the [`JupyterCodeExecutor`](/docs/reference/coding/jupyter/jupyter_code_executor) with the Jupyter server.\n",
    "- Ensure the `jupyter kernelgateway` is started with:\n",
    "    - `--JupyterApp.answer_yes=true` - this ensures that the kernel gateway does not prompt for confirmation when shut down.\n",
    "    - `--JupyterWebsocketPersonality.list_kernels=true` - this ensures that the kernel gateway lists the available kernels.\n",
    "\n",
    "\n",
    "If you wanted to add extra dependencies (for example `matplotlib` and `numpy`) to this image you could customize the Dockerfile like so:\n",
    "\n",
    "```Dockerfile\n",
    "FROM quay.io/jupyter/docker-stacks-foundation\n",
    "\n",
    "SHELL [\"/bin/bash\", \"-o\", \"pipefail\", \"-c\"]\n",
    "\n",
    "USER ${NB_UID}\n",
    "RUN mamba install --yes jupyter_kernel_gateway ipykernel matplotlib numpy &&\n",
    "    mamba clean --all -f -y &&\n",
    "    fix-permissions \"${CONDA_DIR}\" &&\n",
    "    fix-permissions \"/home/${NB_USER}\"\n",
    "\n",
    "ENV TOKEN=\"UNSET\"\n",
    "CMD python -m jupyter kernelgateway \\\n",
    "    --KernelGatewayApp.ip=0.0.0.0 \\\n",
    "    --KernelGatewayApp.port=8888 \\\n",
    "    --KernelGatewayApp.auth_token=\"${TOKEN}\" \\\n",
    "    --JupyterApp.answer_yes=true \\\n",
    "    --JupyterWebsocketPersonality.list_kernels=true\n",
    "\n",
    "EXPOSE 8888\n",
    "\n",
    "WORKDIR \"${HOME}\"\n",
    "```\n",
    "\n",
    "````{=mdx}\n",
    ":::tip\n",
    "To learn about how to combine AutoGen in a Docker image while also executing code in a separate image go [here](/docs/topics/code-execution/cli-code-executor#combining-autogen-in-docker-with-a-docker-based-executor).\n",
    ":::\n",
    "````"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Local\n",
    "\n",
    "````mdx-code-block\n",
    ":::danger\n",
    "The local version will run code on your local system. Use it with caution.\n",
    ":::\n",
    "````\n",
    "\n",
    "To run a local Jupyter server, the [`LocalJupyterServer`](/docs/reference/coding/jupyter/local_jupyter_server#localjupyterserver) can be used.\n",
    "\n",
    "````{=mdx}\n",
    ":::warning\n",
    "The [`LocalJupyterServer`](/docs/reference/coding/jupyter/local_jupyter_server#localjupyterserver) does not function on Windows due to a bug. In this case, you can use the [`DockerJupyterServer`](/docs/reference/coding/jupyter/docker_jupyter_server#dockerjupyterserver) instead or use the [`EmbeddedIPythonCodeExecutor`](/docs/reference/coding/jupyter/embedded_ipython_code_executor). Do note that the intention is to remove the [`EmbeddedIPythonCodeExecutor`](/docs/reference/coding/jupyter/embedded_ipython_code_executor) when the bug is fixed.\n",
    ":::\n",
    "````"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from autogen.coding import CodeBlock\n",
    "from autogen.coding.jupyter import JupyterCodeExecutor, LocalJupyterServer\n",
    "\n",
    "with LocalJupyterServer() as server:\n",
    "    executor = JupyterCodeExecutor(server)\n",
    "    print(\n",
    "        executor.execute_code_blocks(\n",
    "            code_blocks=[\n",
    "                CodeBlock(language=\"python\", code=\"print('Hello, World!')\"),\n",
    "            ]\n",
    "        )\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Remote\n",
    "\n",
    "The [`JupyterCodeExecutor`](/docs/reference/coding/jupyter/jupyter_code_executor) can also connect to a remote Jupyter server. This is done by passing connection information rather than an actual server object into the [`JupyterCodeExecutor`](/docs/reference/coding/jupyter/jupyter_code_executor) constructor.\n",
    "\n",
    "```python\n",
    "from autogen.coding.jupyter import JupyterCodeExecutor, JupyterConnectionInfo\n",
    "\n",
    "executor = JupyterCodeExecutor(\n",
    "    jupyter_server=JupyterConnectionInfo(host='example.com', use_https=True, port=7893, token='mytoken')\n",
    ")\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Image outputs\n",
    "\n",
    "When Jupyter outputs an image, this is saved as a file into the `output_dir` of the [`JupyterCodeExecutor`](/docs/reference/coding/jupyter/jupyter_code_executor), as specified by the constructor. By default this is the current working directory.\n",
    "\n",
    "## Assigning to an agent\n",
    "\n",
    "A single server can support multiple agents, as each executor will create its own kernel. To assign an executor to an agent it can be done like so:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "from autogen import ConversableAgent\n",
    "from autogen.coding.jupyter import DockerJupyterServer, JupyterCodeExecutor\n",
    "\n",
    "server = DockerJupyterServer()\n",
    "\n",
    "output_dir = Path(\"coding\")\n",
    "output_dir.mkdir(exist_ok=True)\n",
    "\n",
    "code_executor_agent = ConversableAgent(\n",
    "    name=\"code_executor_agent\",\n",
    "    llm_config=False,\n",
    "    code_execution_config={\n",
    "        \"executor\": JupyterCodeExecutor(server, output_dir=output_dir),\n",
    "    },\n",
    "    human_input_mode=\"NEVER\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When using code execution it is critical that you update the system prompt of agents you expect to write code to be able to make use of the executor. For example, for the [`JupyterCodeExecutor`](/docs/reference/coding/jupyter/jupyter_code_executor) you might setup a code writing agent like so:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The code writer agent's system message is to instruct the LLM on how to\n",
    "# use the Jupyter code executor with IPython kernel.\n",
    "code_writer_system_message = \"\"\"\n",
    "You have been given coding capability to solve tasks using Python code in a stateful IPython kernel.\n",
    "You are responsible for writing the code, and the user is responsible for executing the code.\n",
    "\n",
    "When you write Python code, put the code in a markdown code block with the language set to Python.\n",
    "For example:\n",
    "```python\n",
    "x = 3\n",
    "```\n",
    "You can use the variable `x` in subsequent code blocks.\n",
    "```python\n",
    "print(x)\n",
    "```\n",
    "\n",
    "Write code incrementally and leverage the statefulness of the kernel to avoid repeating code.\n",
    "Import libraries in a separate code block.\n",
    "Define a function or a class in a separate code block.\n",
    "Run code that produces output in a separate code block.\n",
    "Run code that involves expensive operations like download, upload, and call external APIs in a separate code block.\n",
    "\n",
    "When your code produces an output, the output will be returned to you.\n",
    "Because you have limited conversation memory, if your code creates an image,\n",
    "the output will be a path to the image instead of the image itself.\"\"\"\n",
    "\n",
    "import os\n",
    "\n",
    "code_writer_agent = ConversableAgent(\n",
    "    \"code_writer\",\n",
    "    system_message=code_writer_system_message,\n",
    "    llm_config={\"config_list\": [{\"model\": \"gpt-4\", \"api_key\": os.environ[\"OPENAI_API_KEY\"]}]},\n",
    "    code_execution_config=False,  # Turn off code execution for this agent.\n",
    "    max_consecutive_auto_reply=2,\n",
    "    human_input_mode=\"NEVER\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then we can use these two agents to solve a problem:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[33mcode_executor_agent\u001b[0m (to code_writer):\n",
      "\n",
      "Write Python code to calculate the 14th Fibonacci number.\n",
      "\n",
      "--------------------------------------------------------------------------------\n",
      "\u001b[33mcode_writer\u001b[0m (to code_executor_agent):\n",
      "\n",
      "Sure. The Fibonacci sequence is a series of numbers where the next number is found by adding up the two numbers before it. We know that the first two Fibonacci numbers are 0 and 1. After that, the series looks like:\n",
      "\n",
      "0, 1, 1, 2, 3, 5, 8, 13, 21, 34, 55, ...\n",
      "\n",
      "So, let's define a Python function to calculate the nth Fibonacci number.\n",
      "\n",
      "--------------------------------------------------------------------------------\n",
      "\u001b[33mcode_executor_agent\u001b[0m (to code_writer):\n",
      "\n",
      "\n",
      "\n",
      "--------------------------------------------------------------------------------\n",
      "\u001b[33mcode_writer\u001b[0m (to code_executor_agent):\n",
      "\n",
      "Here is the Python function to calculate the nth Fibonacci number:\n",
      "\n",
      "```python\n",
      "def fibonacci(n):\n",
      "    if n <= 1:\n",
      "        return n\n",
      "    else:\n",
      "        a, b = 0, 1\n",
      "        for _ in range(2, n+1):\n",
      "            a, b = b, a+b\n",
      "        return b\n",
      "```\n",
      "\n",
      "Now, let's use this function to calculate the 14th Fibonacci number.\n",
      "\n",
      "```python\n",
      "fibonacci(14)\n",
      "```\n",
      "\n",
      "--------------------------------------------------------------------------------\n",
      "\u001b[33mcode_executor_agent\u001b[0m (to code_writer):\n",
      "\n",
      "exitcode: 0 (execution succeeded)\n",
      "Code output: \n",
      "377\n",
      "\n",
      "--------------------------------------------------------------------------------\n",
      "ChatResult(chat_id=None,\n",
      "           chat_history=[{'content': 'Write Python code to calculate the 14th '\n",
      "                                     'Fibonacci number.',\n",
      "                          'role': 'assistant'},\n",
      "                         {'content': 'Sure. The Fibonacci sequence is a series '\n",
      "                                     'of numbers where the next number is '\n",
      "                                     'found by adding up the two numbers '\n",
      "                                     'before it. We know that the first two '\n",
      "                                     'Fibonacci numbers are 0 and 1. After '\n",
      "                                     'that, the series looks like:\\n'\n",
      "                                     '\\n'\n",
      "                                     '0, 1, 1, 2, 3, 5, 8, 13, 21, 34, 55, '\n",
      "                                     '...\\n'\n",
      "                                     '\\n'\n",
      "                                     \"So, let's define a Python function to \"\n",
      "                                     'calculate the nth Fibonacci number.',\n",
      "                          'role': 'user'},\n",
      "                         {'content': '', 'role': 'assistant'},\n",
      "                         {'content': 'Here is the Python function to calculate '\n",
      "                                     'the nth Fibonacci number:\\n'\n",
      "                                     '\\n'\n",
      "                                     '```python\\n'\n",
      "                                     'def fibonacci(n):\\n'\n",
      "                                     '    if n <= 1:\\n'\n",
      "                                     '        return n\\n'\n",
      "                                     '    else:\\n'\n",
      "                                     '        a, b = 0, 1\\n'\n",
      "                                     '        for _ in range(2, n+1):\\n'\n",
      "                                     '            a, b = b, a+b\\n'\n",
      "                                     '        return b\\n'\n",
      "                                     '```\\n'\n",
      "                                     '\\n'\n",
      "                                     \"Now, let's use this function to \"\n",
      "                                     'calculate the 14th Fibonacci number.\\n'\n",
      "                                     '\\n'\n",
      "                                     '```python\\n'\n",
      "                                     'fibonacci(14)\\n'\n",
      "                                     '```',\n",
      "                          'role': 'user'},\n",
      "                         {'content': 'exitcode: 0 (execution succeeded)\\n'\n",
      "                                     'Code output: \\n'\n",
      "                                     '377',\n",
      "                          'role': 'assistant'}],\n",
      "           summary='exitcode: 0 (execution succeeded)\\nCode output: \\n377',\n",
      "           cost=({'gpt-4-0613': {'completion_tokens': 193,\n",
      "                                 'cost': 0.028499999999999998,\n",
      "                                 'prompt_tokens': 564,\n",
      "                                 'total_tokens': 757},\n",
      "                  'total_cost': 0.028499999999999998},\n",
      "                 {'gpt-4-0613': {'completion_tokens': 193,\n",
      "                                 'cost': 0.028499999999999998,\n",
      "                                 'prompt_tokens': 564,\n",
      "                                 'total_tokens': 757},\n",
      "                  'total_cost': 0.028499999999999998}),\n",
      "           human_input=[])\n"
     ]
    }
   ],
   "source": [
    "import pprint\n",
    "\n",
    "chat_result = code_executor_agent.initiate_chat(\n",
    "    code_writer_agent, message=\"Write Python code to calculate the 14th Fibonacci number.\"\n",
    ")\n",
    "\n",
    "pprint.pprint(chat_result)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, stop the server. Or better yet use a context manager for it to be stopped automatically."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "server.stop()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "autogen",
   "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.11.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
