ag2 0.10.2__py3-none-any.whl

autogen/_website/notebook_processor.py

@@ -0,0 +1,475 @@
+# Copyright (c) 2023 - 2025, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+
+import argparse
+import concurrent.futures
+import functools
+import json
+import os
+import shutil
+import signal
+import subprocess
+import sys
+import tempfile
+import threading
+import time
+from collections.abc import Callable
+from dataclasses import dataclass
+from functools import lru_cache
+from pathlib import Path
+from typing import Any, TypeVar
+
+from ..import_utils import optional_import_block, require_optional_import
+
+with optional_import_block():
+    import nbformat
+    from nbclient.client import NotebookClient
+    from nbclient.exceptions import (
+        CellExecutionError,
+        CellTimeoutError,
+    )
+    from nbformat import NotebookNode
+    from termcolor import colored
+
+
+# Notebook execution based on nbmake: https://github.com/treebeardtech/nbmakes
+@dataclass
+class NotebookError:
+    error_name: str
+    error_value: str | None
+    traceback: str
+    cell_source: str
+
+
+@dataclass
+class NotebookSkip:
+    reason: str
+
+
+NB_VERSION = 4
+
+
+class Result:
+    def __init__(self, returncode: int, stdout: str, stderr: str):
+        self.returncode = returncode
+        self.stdout = stdout
+        self.stderr = stderr
+
+
+def path(path_str: str) -> Path:
+    """Return a Path object."""
+    return Path(path_str)
+
+
+@lru_cache
+def check_quarto_bin(quarto_bin: str = "quarto") -> bool:
+    """Check if quarto is installed."""
+    try:
+        version_str = subprocess.check_output([quarto_bin, "--version"], text=True).strip()
+        version = tuple(map(int, version_str.split(".")))
+        return version >= (1, 5, 23)
+
+    except FileNotFoundError:
+        return False
+
+
+C = TypeVar("C", bound=Callable[..., Any])
+
+
+def require_quarto_bin(f: C) -> C:
+    """Decorator to skip a function if quarto is not installed."""
+    if check_quarto_bin():
+        return f
+    else:
+
+        @functools.wraps(f)
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
+            return ImportError("Quarto is not installed")
+
+        return wrapper  # type: ignore[return-value]
+
+
+def load_metadata(notebook: Path) -> dict[str, dict[str, str | list[str] | None]]:
+    content = json.load(notebook.open(encoding="utf-8"))
+    metadata: dict[str, dict[str, str | list[str] | None]] = content.get("metadata", {})
+    return metadata
+
+
+def skip_reason_or_none_if_ok(notebook: Path) -> str | None | dict[str, Any]:
+    """Return a reason to skip the notebook, or None if it should not be skipped."""
+    if notebook.suffix != ".ipynb":
+        return "not a notebook"
+
+    if not notebook.exists():
+        return "file does not exist"
+
+    # Extra checks for notebooks in the notebook directory
+    if "notebook" not in notebook.parts:
+        return None
+
+    with open(notebook, encoding="utf-8") as f:
+        content = f.read()
+
+    # Load the json and get the first cell
+    json_content = json.loads(content)
+    first_cell = json_content["cells"][0]
+
+    # <!-- and --> must exists on lines on their own
+    if first_cell["cell_type"] == "markdown" and first_cell["source"][0].strip() == "<!--":
+        raise ValueError(
+            f"Error in {notebook.resolve()!s} - Front matter should be defined in the notebook metadata now."
+        )
+
+    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 ⚠️"
+
+    front_matter = metadata["front_matter"]
+
+    if "tags" not in front_matter:
+        return "tags is not in front matter"
+
+    if "description" not in front_matter:
+        return "description is not in front matter"
+
+    # Make sure tags is a list of strings
+    if front_matter["tags"] is not None and not all(isinstance(tag, str) for tag in front_matter["tags"]):
+        return "tags must be a list of strings"
+
+    # Make sure description is a string
+    if not isinstance(front_matter["description"], str):
+        return "description must be a string"
+
+    return None
+
+
+def extract_title(notebook: Path) -> str | None:
+    """Extract the title of the notebook."""
+    with open(notebook, encoding="utf-8") as f:
+        content = f.read()
+
+    # Load the json and get the first cell
+    json_content = json.loads(content)
+    first_cell = json_content["cells"][0]
+
+    # find the # title
+    for line in first_cell["source"]:
+        if line.startswith("# "):
+            title: str = line[2:].strip()
+            # Strip off the { if it exists
+            if "{" in title:
+                title = title[: title.find("{")].strip()
+            return title
+
+    return None
+
+
+def start_thread_to_terminate_when_parent_process_dies(ppid: int) -> None:
+    pid = os.getpid()
+
+    def f() -> None:
+        while True:
+            try:
+                os.kill(ppid, 0)
+            except OSError:
+                os.kill(pid, signal.SIGTERM)
+            time.sleep(1)
+
+    thread = threading.Thread(target=f, daemon=True)
+    thread.start()
+
+
+@require_optional_import("termcolor", "docs")
+def fmt_skip(notebook: Path, reason: str) -> str:
+    return f"{colored('[Skip]', 'yellow')} {colored(notebook.name, 'blue')}: {reason}"
+
+
+@require_optional_import("termcolor", "docs")
+def fmt_ok(notebook: Path) -> str:
+    return f"{colored('[OK]', 'green')} {colored(notebook.name, 'blue')} ✅"
+
+
+@require_optional_import("termcolor", "docs")
+def fmt_error(notebook: Path, error: 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")
+
+
+@require_quarto_bin
+@require_optional_import("nbclient", "docs")
+def test_notebook(notebook_path: Path, timeout: int = 300) -> tuple[Path, NotebookError | NotebookSkip | None]:
+    nb = nbformat.read(str(notebook_path), NB_VERSION)  # type: ignore[arg-type,no-untyped-call]
+
+    if "skip_test" in nb.metadata:
+        return notebook_path, NotebookSkip(reason=nb.metadata.skip_test)
+
+    try:
+        c = NotebookClient(
+            nb,
+            timeout=timeout,
+            allow_errors=False,
+            record_timing=True,
+        )
+        os.environ["PYDEVD_DISABLE_FILE_VALIDATION"] = "1"
+        os.environ["TOKENIZERS_PARALLELISM"] = "false"
+        with tempfile.TemporaryDirectory() as tempdir:
+            c.execute(cwd=tempdir)
+    except CellExecutionError:
+        error = get_error_info(nb)
+        assert error is not None
+        return notebook_path, error
+    except CellTimeoutError:
+        error = get_timeout_info(nb)
+        assert error is not None
+        return notebook_path, error
+
+    return notebook_path, None
+
+
+# Find the first code cell which did not complete.
+@require_optional_import("nbclient", "docs")
+def get_timeout_info(
+    nb: NotebookNode,
+) -> NotebookError | None:
+    for i, cell in enumerate(nb.cells):
+        if cell.cell_type != "code":
+            continue
+        if "shell.execute_reply" not in cell.metadata.execution:
+            return NotebookError(
+                error_name="timeout",
+                error_value="",
+                traceback="",
+                cell_source="".join(cell["source"]),
+            )
+
+    return None
+
+
+@require_optional_import("nbclient", "docs")
+def get_error_info(nb: NotebookNode) -> NotebookError | None:
+    for cell in nb["cells"]:  # get LAST error
+        if cell["cell_type"] != "code":
+            continue
+        errors = [output for output in cell["outputs"] if output["output_type"] == "error" or "ename" in output]
+
+        if errors:
+            traceback = "\n".join(errors[0].get("traceback", ""))
+            return NotebookError(
+                error_name=errors[0].get("ename", ""),
+                error_value=errors[0].get("evalue", ""),
+                traceback=traceback,
+                cell_source="".join(cell["source"]),
+            )
+    return None
+
+
+def collect_notebooks(notebook_directory: Path, website_build_directory: Path) -> list[Path]:
+    notebooks = list(notebook_directory.glob("*.ipynb"))
+    notebooks.extend(list(website_build_directory.glob("docs/**/*.ipynb")))
+    return notebooks
+
+
+@require_quarto_bin
+@require_optional_import(["nbclient", "termcolor"], "docs")
+def process_notebook(
+    src_notebook: Path,
+    website_build_directory: Path,
+    notebook_dir: Path,
+    quarto_bin: str,
+    dry_run: bool,
+    target_dir_func: Callable[[Path], Path],
+    post_processor: Callable[[Path, Path, dict[str, Any], Path], None] | None = None,
+) -> str:
+    """Process a single notebook.
+
+    Args:
+        src_notebook: Source notebook path
+        website_build_directory: Output directory
+        notebook_dir: Base notebooks directory
+        quarto_bin: Path to quarto binary
+        dry_run: If True, don't actually process
+        target_dir_func: Function to determine target directory for notebooks
+        post_processor: Optional callback for post-processing
+    """
+    in_notebook_dir = "notebook" in src_notebook.parts
+
+    metadata = load_metadata(src_notebook)
+
+    title = extract_title(src_notebook)
+    if title is None:
+        return fmt_error(src_notebook, "Title not found in notebook")
+
+    front_matter = {}
+    if "front_matter" in metadata:
+        front_matter = metadata["front_matter"]
+
+    front_matter["title"] = title
+
+    if in_notebook_dir:
+        relative_notebook = src_notebook.resolve().relative_to(notebook_dir.resolve())
+        dest_dir = target_dir_func(website_build_directory)
+        target_file = dest_dir / relative_notebook.with_suffix(".mdx")
+        intermediate_notebook = dest_dir / relative_notebook
+
+        # If the intermediate_notebook already exists, check if it is newer than the source file
+        if target_file.exists() and target_file.stat().st_mtime > src_notebook.stat().st_mtime:
+            return fmt_skip(src_notebook, f"target file ({target_file.name}) is newer ☑️")
+
+        if dry_run:
+            return colored(f"Would process {src_notebook.name}", "green")
+
+        # Copy notebook to target dir
+        # The reason we copy the notebook is that quarto does not support rendering from a different directory
+        shutil.copy(src_notebook, intermediate_notebook)
+
+        # Check if another file has to be copied too
+        # Solely added for the purpose of agent_library_example.json
+        if "extra_files_to_copy" in metadata:
+            for file in metadata["extra_files_to_copy"]:
+                shutil.copy(src_notebook.parent / file, dest_dir / file)
+
+        # Capture output
+        result = subprocess.run([quarto_bin, "render", intermediate_notebook], capture_output=True, text=True)
+        if result.returncode != 0:
+            return fmt_error(
+                src_notebook, f"Failed to render {src_notebook}\n\nstderr:\n{result.stderr}\nstdout:\n{result.stdout}"
+            )
+
+        # Unlink intermediate files
+        intermediate_notebook.unlink()
+    else:
+        target_file = src_notebook.with_suffix(".mdx")
+
+        # If the intermediate_notebook already exists, check if it is newer than the source file
+        if target_file.exists() and target_file.stat().st_mtime > src_notebook.stat().st_mtime:
+            return fmt_skip(src_notebook, f"target file ({target_file.name}) is newer ☑️")
+
+        if dry_run:
+            return colored(f"Would process {src_notebook.name}", "green")
+
+        result = subprocess.run([quarto_bin, "render", src_notebook], capture_output=True, text=True)
+        if result.returncode != 0:
+            return fmt_error(
+                src_notebook, f"Failed to render {src_notebook}\n\nstderr:\n{result.stderr}\nstdout:\n{result.stdout}"
+            )
+
+    # Use post-processor if provided
+    if post_processor and not dry_run:
+        post_processor(target_file, src_notebook, front_matter, website_build_directory)
+
+    return fmt_ok(src_notebook)
+
+
+def create_base_argument_parser() -> argparse.ArgumentParser:
+    """Create the base argument parser with common options."""
+    parser = argparse.ArgumentParser()
+    subparsers = parser.add_subparsers(dest="subcommand")
+
+    parser.add_argument(
+        "--notebook-directory",
+        type=path,
+        help="Directory containing notebooks to process",
+    )
+    parser.add_argument("--website-build-directory", type=path, help="Root directory of website build")
+    parser.add_argument("--force", help="Force re-rendering of all notebooks", action="store_true", default=False)
+
+    render_parser = subparsers.add_parser("render")
+    render_parser.add_argument("--quarto-bin", help="Path to quarto binary", default="quarto")
+    render_parser.add_argument("--dry-run", help="Don't render", action="store_true")
+    render_parser.add_argument("notebooks", type=path, nargs="*", default=None)
+
+    test_parser = subparsers.add_parser("test")
+    test_parser.add_argument("--timeout", help="Timeout for each notebook", type=int, default=60)
+    test_parser.add_argument("--exit-on-first-fail", "-e", help="Exit after first test fail", action="store_true")
+    test_parser.add_argument("notebooks", type=path, nargs="*", default=None)
+    test_parser.add_argument("--workers", help="Number of workers to use", type=int, default=-1)
+
+    return parser
+
+
+def process_notebooks_core(
+    args: argparse.Namespace,
+    post_process_func: Callable[[Path, Path, dict[str, Any], Path], None] | None,
+    target_dir_func: Callable[[Path], Path],
+) -> list[Path]:
+    """Core logic for processing notebooks shared across build systems.
+
+    Args:
+        args: Command line arguments
+        post_process_func: Function for post-processing rendered notebooks
+        target_dir_func: Function to determine target directory for notebooks
+    """
+    collected_notebooks = (
+        args.notebooks if args.notebooks else collect_notebooks(args.notebook_directory, args.website_build_directory)
+    )
+
+    filtered_notebooks = []
+    for notebook in collected_notebooks:
+        reason = skip_reason_or_none_if_ok(notebook)
+        if reason and isinstance(reason, str):
+            print(fmt_skip(notebook, reason))
+        else:
+            filtered_notebooks.append(notebook)
+
+    if args.subcommand == "test":
+        if args.workers == -1:
+            args.workers = None
+        failure = False
+        with concurrent.futures.ProcessPoolExecutor(
+            max_workers=args.workers,
+            initializer=start_thread_to_terminate_when_parent_process_dies,
+            initargs=(os.getpid(),),
+        ) as executor:
+            futures = [executor.submit(test_notebook, f, args.timeout) for f in filtered_notebooks]
+            for future in concurrent.futures.as_completed(futures):
+                notebook, optional_error_or_skip = future.result()
+                if isinstance(optional_error_or_skip, NotebookError):
+                    if optional_error_or_skip.error_name == "timeout":
+                        print(fmt_error(notebook, optional_error_or_skip.error_name))
+                    else:
+                        print("-" * 80)
+                        print(fmt_error(notebook, optional_error_or_skip))
+                        print(optional_error_or_skip.traceback)
+                        print("-" * 80)
+                    if args.exit_on_first_fail:
+                        sys.exit(1)
+                    failure = True
+                elif isinstance(optional_error_or_skip, NotebookSkip):
+                    print(fmt_skip(notebook, optional_error_or_skip.reason))
+                else:
+                    print(fmt_ok(notebook))
+
+        if failure:
+            sys.exit(1)
+
+    elif args.subcommand == "render":
+        check_quarto_bin(args.quarto_bin)
+
+        target_dir = target_dir_func(args.website_build_directory)
+        if not target_dir.exists():
+            target_dir.mkdir(parents=True)
+
+        for notebook in filtered_notebooks:
+            print(
+                process_notebook(
+                    notebook,
+                    args.website_build_directory,
+                    args.notebook_directory,
+                    args.quarto_bin,
+                    args.dry_run,
+                    target_dir_func,
+                    post_process_func,
+                )
+            )
+
+    return filtered_notebooks