diff --git a/autogen/coding/jupyter/base.py b/autogen/coding/jupyter/base.py index 8e86897249e..d896b6ac3cc 100644 --- a/autogen/coding/jupyter/base.py +++ b/autogen/coding/jupyter/base.py @@ -7,9 +7,13 @@ class JupyterConnectionInfo: """(Experimental)""" host: str + """`str` - Host of the Jupyter gateway server""" use_https: bool + """`bool` - Whether to use HTTPS""" port: int + """`int` - Port of the Jupyter gateway server""" token: Optional[str] + """`Optional[str]` - Token for authentication. If None, no token is used""" @runtime_checkable @@ -18,4 +22,5 @@ class JupyterConnectable(Protocol): @property def connection_info(self) -> JupyterConnectionInfo: + """Return the connection information for this connectable.""" pass diff --git a/autogen/coding/jupyter/jupyter_client.py b/autogen/coding/jupyter/jupyter_client.py index 459add85b5f..e02d3b49bca 100644 --- a/autogen/coding/jupyter/jupyter_client.py +++ b/autogen/coding/jupyter/jupyter_client.py @@ -23,9 +23,12 @@ class JupyterClient: - """(Experimental) A client for communicating with a Jupyter gateway server.""" - def __init__(self, connection_info: JupyterConnectionInfo): + """(Experimental) A client for communicating with a Jupyter gateway server. + + Args: + connection_info (JupyterConnectionInfo): Connection information + """ self._connection_info = connection_info self._session = requests.Session() retries = Retry(total=5, backoff_factor=0.1) diff --git a/autogen/coding/jupyter/jupyter_code_executor.py b/autogen/coding/jupyter/jupyter_code_executor.py index 5e190d5f1b9..6e63652ef37 100644 --- a/autogen/coding/jupyter/jupyter_code_executor.py +++ b/autogen/coding/jupyter/jupyter_code_executor.py @@ -22,27 +22,6 @@ class JupyterCodeExecutor(CodeExecutor): - """(Experimental) A code executor class that executes code statefully using an embedded - IPython kernel managed by this class. - - **This will execute LLM generated code on the local machine.** - - Each execution is stateful and can access variables created from previous - executions in the same session. The kernel must be installed before using - this class. The kernel can be installed using the following command: - `python -m ipykernel install --user --name {kernel_name}` - where `kernel_name` is the name of the kernel to install. - - Args: - timeout (int): The timeout for code execution, by default 60. - kernel_name (str): The kernel name to use. Make sure it is installed. - By default, it is "python3". - output_dir (str): The directory to save output files, by default ".". - system_message_update (str): The system message update to add to the - agent that produces code. By default it is - `JupyterCodeExecutor.DEFAULT_SYSTEM_MESSAGE_UPDATE`. - """ - DEFAULT_SYSTEM_MESSAGE_UPDATE: ClassVar[ str ] = """ @@ -72,24 +51,10 @@ class JupyterCodeExecutor(CodeExecutor): """ class UserCapability: - """(Experimental) An AgentCapability class that gives agent ability use a stateful - IPython code executor. This capability can be added to an agent using - the `add_to_agent` method which append a system message update to the - agent's system message.""" - def __init__(self, system_message_update: str): self._system_message_update = system_message_update def add_to_agent(self, agent: LLMAgent) -> None: - """Add this capability to an agent by appending a system message - update to the agent's system message. - - **Currently we do not check for conflicts with existing content in - the agent's system message.** - - Args: - agent (LLMAgent): The agent to add the capability to. - """ agent.update_system_message(agent.system_message + self._system_message_update) def __init__( @@ -100,6 +65,19 @@ def __init__( output_dir: Union[Path, str] = Path("."), system_message_update: str = DEFAULT_SYSTEM_MESSAGE_UPDATE, ): + """(Experimental) A code executor class that executes code statefully using + a Jupyter server supplied to this class. + + Each execution is stateful and can access variables created from previous + executions in the same session. + + Args: + jupyter_server (Union[JupyterConnectable, JupyterConnectionInfo]): The Jupyter server to use. + timeout (int): The timeout for code execution, by default 60. + kernel_name (str): The kernel name to use. Make sure it is installed. + By default, it is "python3". + output_dir (str): The directory to save output files, by default ".". + """ if timeout < 1: raise ValueError("Timeout must be greater than or equal to 1.") @@ -130,8 +108,6 @@ def __init__( @property def user_capability(self) -> "JupyterCodeExecutor.UserCapability": - """(Experimental) Export a user capability for this executor that can be added to - an agent using the `add_to_agent` method.""" return JupyterCodeExecutor.UserCapability(self._system_message_update) @property @@ -142,8 +118,7 @@ def code_extractor(self) -> CodeExtractor: def execute_code_blocks(self, code_blocks: List[CodeBlock]) -> IPythonCodeResult: """(Experimental) Execute a list of code blocks and return the result. - This method executes a list of code blocks as cells in an IPython kernel - managed by this class. + This method executes a list of code blocks as cells in the Jupyter kernel. See: https://jupyter-client.readthedocs.io/en/stable/messaging.html for the message protocol. diff --git a/autogen/coding/jupyter/local_jupyter_server.py b/autogen/coding/jupyter/local_jupyter_server.py index decbb3f430e..5ad0fd01a52 100644 --- a/autogen/coding/jupyter/local_jupyter_server.py +++ b/autogen/coding/jupyter/local_jupyter_server.py @@ -39,6 +39,17 @@ def __init__( log_max_bytes: int = 1048576, log_backup_count: int = 3, ): + """Runs a Jupyter Kernel Gateway server locally. + + Args: + ip (str, optional): IP address to bind to. Defaults to "127.0.0.1". + port (Optional[int], optional): Port to use, if None it automatically selects a port. Defaults to None. + token (Union[str, GenerateToken], optional): Token to use for Jupyter server. By default will generate a token. Using None will use no token for authentication. Defaults to GenerateToken(). + log_file (str, optional): File for Jupyter Kernel Gateway logs. Defaults to "jupyter_gateway.log". + log_level (str, optional): Level for Jupyter Kernel Gateway logs. Defaults to "INFO". + log_max_bytes (int, optional): Max logfile size. Defaults to 1048576. + log_backup_count (int, optional): Number of backups for rotating log. Defaults to 3. + """ # Remove as soon as https://github.com/jupyter-server/kernel_gateway/issues/398 is fixed if sys.platform == "win32": raise ValueError("LocalJupyterServer is not supported on Windows due to kernelgateway bug.") diff --git a/website/.gitignore b/website/.gitignore index be132bdc827..9d2fc91aa03 100644 --- a/website/.gitignore +++ b/website/.gitignore @@ -12,6 +12,7 @@ docs/reference /docs/notebooks docs/topics/llm_configuration.mdx +docs/topics/code-execution/jupyter-code-executor.mdx # Misc .DS_Store diff --git a/website/docs/topics/code-execution/_category_.json b/website/docs/topics/code-execution/_category_.json new file mode 100644 index 00000000000..57f54268afb --- /dev/null +++ b/website/docs/topics/code-execution/_category_.json @@ -0,0 +1,5 @@ +{ + "position": 2, + "label": "Code Execution", + "collapsible": true +} diff --git a/website/docs/topics/code-execution/jupyter-code-executor.ipynb b/website/docs/topics/code-execution/jupyter-code-executor.ipynb new file mode 100644 index 00000000000..1e1e48ad479 --- /dev/null +++ b/website/docs/topics/code-execution/jupyter-code-executor.ipynb @@ -0,0 +1,454 @@ +{ + "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 local, in Docker or even a remove 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 'pyautogen[jupyter-executor]'\n", + "```\n", + "\n", + "## Jupyter Server\n", + "\n", + "### Local\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 [`EmbeddedJupyterServer`](/docs/reference/coding/jupyter/embedded_ipython_code_executor). Do note that the intention is to remove the [`EmbeddedJupyterServer`](/docs/reference/coding/jupyter/embedded_ipython_code_executor) when the bug is fixed.\n", + ":::\n", + "````" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "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 LocalJupyterServer, JupyterCodeExecutor\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": [ + "### 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" + ] + }, + { + "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 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 autogen import ConversableAgent\n", + "from autogen.coding.jupyter import DockerJupyterServer, JupyterCodeExecutor\n", + "from pathlib import Path\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.7" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/website/process_notebooks.py b/website/process_notebooks.py index 2246e07bbd9..8c4651fdc66 100755 --- a/website/process_notebooks.py +++ b/website/process_notebooks.py @@ -105,6 +105,9 @@ def skip_reason_or_none_if_ok(notebook: Path) -> typing.Optional[str]: metadata = load_metadata(notebook) + if "skip_render" in metadata: + return metadata["skip_render"] + if "front_matter" not in metadata: return "front matter missing from notebook metadata ⚠️" @@ -154,8 +157,6 @@ def process_notebook(src_notebook: Path, website_dir: Path, notebook_dir: Path, in_notebook_dir = "notebook" in src_notebook.parts metadata = load_metadata(src_notebook) - if "skip_render" in metadata: - return fmt_skip(src_notebook, "skip_render is in notebook metadata") title = extract_title(src_notebook) if title is None: @@ -331,7 +332,7 @@ def post_process_mdx(rendered_mdx: Path, source_notebooks: Path, front_matter: D else: title_search_content = content - title_exists = title_search_content.find("# ") != -1 + title_exists = title_search_content.find("\n# ") != -1 if not title_exists: content = f"# {front_matter['title']}\n{content}" @@ -389,16 +390,21 @@ def collect_notebooks(notebook_directory: Path, website_directory: Path) -> typi return notebooks -def fmt_skip(notebook: Path, reason: str) -> None: +def fmt_skip(notebook: Path, reason: str) -> str: return f"{colored('[Skip]', 'yellow')} {colored(notebook.name, 'blue')}: {reason}" -def fmt_ok(notebook: Path) -> None: +def fmt_ok(notebook: Path) -> str: return f"{colored('[OK]', 'green')} {colored(notebook.name, 'blue')} ✅" -def fmt_error(notebook: Path, error: NotebookError) -> None: - return f"{colored('[Error]', 'red')} {colored(notebook.name, 'blue')}: {error.error_name} - {error.error_value}" +def fmt_error(notebook: Path, error: Union[NotebookError, str]) -> str: + if isinstance(error, str): + return f"{colored('[Error]', 'red')} {colored(notebook.name, 'blue')}: {error}" + elif isinstance(error, NotebookError): + return f"{colored('[Error]', 'red')} {colored(notebook.name, 'blue')}: {error.error_name} - {error.error_value}" + else: + raise ValueError("error must be a string or a NotebookError") def start_thread_to_terminate_when_parent_process_dies(ppid: int): @@ -480,11 +486,7 @@ def main() -> None: else: print("-" * 80) - print( - fmt_error( - notebook, f"{optional_error_or_skip.error_name} - {optional_error_or_skip.error_value}" - ) - ) + print(fmt_error(notebook, optional_error_or_skip)) print(optional_error_or_skip.traceback) print("-" * 80) if args.exit_on_first_fail: