docs: vcr only set up for markdown (#3339)

- PR branch for testing w/ typescript (@benjamincburns )
- implementation needs better cache invalidation for the cassettes
(@eyurtsev)

---------

Co-authored-by: Ben Burns <803016+benjamincburns@users.noreply.github.com>

docs/_scripts/assets/nock_setup_preamble.ts

@@ -0,0 +1,75 @@
+import nock, { Definition } from "nock";
+import msgpack from "msgpack-lite";
+import zlib from "node:zlib";
+import fs from "node:fs/promises";
+import { Buffer } from "node:buffer";
+
+// deno style imports here because we're running this in the deno jupyter kernel
+
+interface NockCassetteData {
+  hash: string;
+  entries: Definition[];
+}
+
+// Utility functions for compression & serialization
+function compressData(data: NockCassetteData, compressionLevel = 9): string {
+  const packed = msgpack.encode(data);
+  const compressed = zlib.deflateSync(packed, { level: compressionLevel });
+  return compressed.toString("base64");
+}
+
+function decompressData(compressedString: string): NockCassetteData {
+  const decoded = Buffer.from(compressedString, "base64");
+  const decompressed = zlib.inflateSync(decoded);
+  return msgpack.decode(decompressed) as NockCassetteData;
+}
+
+// deno-lint-ignore no-unused-vars
+class HashedCassette {
+  private recording = true;
+
+  constructor(
+    private readonly cassettePath: string,
+    private readonly hash: string
+  ) {}
+
+  async enter() {
+    try {
+      const rawCassette = await fs.readFile(this.cassettePath, "utf-8");
+      const data = decompressData(rawCassette);
+      if (data.hash === this.hash) {
+        this.recording = false;
+        nock.disableNetConnect();
+        nock.define(data.entries);
+        return;
+      }
+    } catch (error) {
+      if (error instanceof Error && error.message.includes("ENOENT")) {
+        this.recording = true;
+      } else {
+        throw error;
+      }
+    }
+
+    nock.recorder.rec({
+      dont_print: true,
+      output_objects: true,
+    });
+  }
+
+  async exit() {
+    if (this.recording) {
+      const entries = nock.recorder.play() as Definition[];
+      const data = {
+        hash: this.hash,
+        entries,
+      };
+      const compressed = compressData(data);
+      await fs.writeFile(this.cassettePath, compressed);
+    } else {
+      nock.enableNetConnect();
+      nock.restore();
+      nock.cleanAll();
+    }
+  }
+}

docs/_scripts/assets/vcr_setup_preamble.py

@@ -0,0 +1,108 @@
+import base64
+import os
+import zlib
+from logging import getLogger
+from types import TracebackType
+from typing import Optional, Any, Type
+
+import msgpack
+import vcr
+
+logger = getLogger(__name__)
+
+os.environ.pop("LANGCHAIN_TRACING_V2", None)
+custom_vcr = vcr.VCR()
+
+
+def compress_data(data: Any, compression_level: int = 9) -> str:
+    packed = msgpack.packb(data, use_bin_type=True)
+    compressed = zlib.compress(packed, level=compression_level)
+    return base64.b64encode(compressed).decode("utf-8")
+
+
+def decompress_data(compressed_string: str) -> Any:
+    decoded = base64.b64decode(compressed_string)
+    decompressed = zlib.decompress(decoded)
+    return msgpack.unpackb(decompressed, raw=False)
+
+
+class AdvancedCompressedSerializer:
+    def serialize(self, cassette_dict: Any) -> str:
+        return compress_data(cassette_dict)
+
+    def deserialize(self, cassette_string: str) -> Any:
+        return decompress_data(cassette_string)
+
+
+custom_vcr.register_serializer("advanced_compressed", AdvancedCompressedSerializer())
+custom_vcr.serializer = "advanced_compressed"
+
+
+class HashedCassette:
+    def __init__(self, cassette_path: str, hash_value: str) -> None:
+        """A context manager for using VCR cassettes with an embedded hash value.
+
+        Args:
+            cassette_path (str): The file path of the cassette (independent of hash).
+            hash_value (str): The expected hash value (e.g. a uuid string).
+
+        This class provides a context manager for using VCR cassettes with an embedded hash value.
+        The hash value is used to ensure that the cassette matches the expected state, and if not,
+        the cassette is removed or updated with the new hash value.
+        """
+        self.cassette_path: str = cassette_path
+        self.hash_value: str = hash_value
+        self.vcr: vcr.VCR = custom_vcr
+        self.cassette_context: Optional[Any] = None
+
+    def __enter__(self) -> Any:
+        # Get the serializer instance from the VCR instance.
+        serializer = self.vcr.serializers[self.vcr.serializer]
+        # If the cassette file exists, check its embedded hash.
+        if os.path.exists(self.cassette_path):
+            with open(self.cassette_path, "r") as f:
+                content = f.read()
+            try:
+                cassette_data = serializer.deserialize(content)
+            except Exception as e:
+                print(f"Error deserializing cassette, removing file: {e}")
+                os.remove(self.cassette_path)
+            else:
+                existing_hash = cassette_data.get("cassette_hash")
+                if existing_hash != self.hash_value:
+                    print("Hash mismatch. Removing outdated cassette.")
+                    os.remove(self.cassette_path)
+        # Now enter the VCR cassette context.
+        self.cassette_context = custom_vcr.use_cassette(
+            self.cassette_path,
+            filter_headers=["x-api-key", "authorization"],
+            record_mode="once",
+            serializer="advanced_compressed",
+        )
+        return self.cassette_context.__enter__()
+
+    def __exit__(
+        self,
+        exc_type: Optional[Type[BaseException]] = None,
+        exc_val: Optional[BaseException] = None,
+        exc_tb: Optional[TracebackType] = None,
+    ) -> Optional[bool]:
+        # Exit the VCR cassette context.
+        result = self.cassette_context.__exit__(exc_type, exc_val, exc_tb)
+        serializer = self.vcr.serializers[self.vcr.serializer]
+        # If a cassette was recorded (or updated), open and update its hash.
+        if os.path.exists(self.cassette_path):
+            with open(self.cassette_path, "r") as f:
+                content = f.read()
+            try:
+                cassette_data = serializer.deserialize(content)
+            except Exception as e:
+                logger.error(f"Error deserializing cassette during exit: {e}")
+                return result
+            # Update the cassette data with the expected hash.
+            if cassette_data.get("cassette_hash") != self.hash_value:
+                cassette_data["cassette_hash"] = self.hash_value
+                serialized_data = serializer.serialize(cassette_data)
+                with open(self.cassette_path, "w") as f:
+                    f.write(serialized_data)
+        return result

docs/_scripts/notebook_hooks.py

@@ -1,14 +1,20 @@
 import logging
 import os
 import re
-from typing import Any, Dict
+import traceback
+from typing import Any, Callable, Dict
 
+from markdown import Markdown
+from pymdownx.superfences import SuperFencesException
 from mkdocs.structure.files import Files, File
 from mkdocs.structure.pages import Page
 import posixpath
 
+from markdown_exec.hooks import SessionHistoryEntry
+
 from generate_api_reference_links import update_markdown_with_imports
 from notebook_convert import convert_notebook
+from setup_vcr import load_postamble, load_preamble, _hash_string
 
 logger = logging.getLogger(__name__)
 logging.basicConfig()
@@ -57,6 +63,29 @@ def on_files(files: Files, **kwargs: Dict[str, Any]):
     return new_files
 
 
+def _add_path_to_code_blocks(markdown: str, page: Page) -> str:
+    """Add the path to the code blocks."""
+    code_block_pattern = re.compile(
+        r"(?P<indent>[ \t]*)```(?P<language>\w+)[ ]*(?P<attributes>[^\n]*)\n"
+        r"(?P<code>((?:.*\n)*?))"  # Capture the code inside the block using named group
+        r"(?P=indent)```"  # Match closing backticks with the same indentation
+    )
+
+    def replace_code_block_header(match: re.Match) -> str:
+        indent = match.group("indent")
+        language = match.group("language")
+        attributes = match.group("attributes").rstrip()
+
+        if 'exec="on"' not in attributes:
+            # Return original code block
+            return match.group(0)
+
+        code = match.group("code")
+        return f'{indent}```{language} {attributes} path="{page.file.src_path}"\n{code}{indent}```'
+
+    return code_block_pattern.sub(replace_code_block_header, markdown)
+
+
 def _highlight_code_blocks(markdown: str) -> str:
     """Find code blocks with highlight comments and add hl_lines attribute.
 
@@ -109,8 +138,8 @@ def _highlight_code_blocks(markdown: str) -> str:
             return (
                 f'{indent}```{language} hl_lines="{" ".join(highlighted_lines)}"\n'
                 # The indent and terminating \n is already included in the code block
-                f'{new_code_block}'
-                f'{indent}```'
+                f"{new_code_block}"
+                f"{indent}```"
             )
         else:
             return (
@@ -125,6 +154,90 @@ def _highlight_code_blocks(markdown: str) -> str:
     return markdown
 
 
+def handle_vcr_setup(
+    *,
+    formatter: Callable,
+    language: str,
+    code: str,
+    session: str,
+    id: str,
+    md: Markdown,
+    **kwargs: Dict[str, Any],
+) -> str:
+    """Handle VCR setup in markdown content if necessary."""
+    try:
+        if kwargs.get("extra", None) is None:
+            raise ValueError(
+                f"error while processing {language} block: extra dict is required"
+            )
+
+        if kwargs["extra"].get("path", None) is None:
+            raise ValueError(
+                f"error while processing {language} block: path is required"
+            )
+
+        document_filename = kwargs["extra"]["path"]
+
+        logger.info("document_filename: %s", document_filename)
+
+        if session is None or session == "" and id is None or id == "":
+            id = _hash_string(code)
+
+        cassette_prefix = document_filename.replace(".md", "").replace(os.path.sep, "_")
+
+        cassette_dir = os.path.abspath(
+            os.path.join(os.path.dirname(os.path.dirname(__file__)), "cassettes")
+        )
+        os.makedirs(cassette_dir, exist_ok=True)
+
+        # Build a unique cassette name.
+        cassette_name = os.path.join(
+            cassette_dir,
+            f"{cassette_prefix}_{session if session else id}_{language}.msgpack.zlib",
+        )
+
+        # Add context manager at start with explicit __enter__ and __exit__ calls
+
+        wrapped_lines = [
+            load_preamble(language, code, cassette_name),
+            code,
+        ]
+
+        if session is None or session == "":
+            wrapped_lines.append(load_postamble(language))
+
+        transformed_source = "\n".join(wrapped_lines)
+        return dict(
+            transform_source=lambda code: (transformed_source, code),
+            id=id,
+            extra={},
+        )
+    except Exception as e:
+        raise SuperFencesException(traceback.format_exc()) from e
+
+
+def handle_vcr_teardown(
+    *,
+    formatter: Callable,
+    language: str,
+    session: str,
+    history: list[SessionHistoryEntry],
+):
+    session = history[-1].inputs["session"]
+    inputs = dict(history[-1].inputs)
+    del inputs["session"]
+    del inputs["code"]
+    del inputs["language"]
+    del inputs["id"]
+    formatter(
+        code="_cassette.__exit__() # markdown-exec: hide",
+        language="python",
+        session=session,
+        id=f"{id}_vcr_end",
+        **inputs,
+    )
+
+
 def _on_page_markdown_with_config(
     markdown: str,
     page: Page,
@@ -135,6 +248,7 @@ def _on_page_markdown_with_config(
 ) -> str:
     if DISABLED:
         return markdown
+
     if page.file.src_path.endswith(".ipynb"):
         logger.info("Processing Jupyter notebook: %s", page.file.src_path)
         markdown = convert_notebook(page.file.abs_src_path)
@@ -145,6 +259,11 @@ def _on_page_markdown_with_config(
     # Apply highlight comments to code blocks
     markdown = _highlight_code_blocks(markdown)
 
+    # Add file path as an attribute to code blocks that are executable.
+    # This file path is used to associate fixtures with the executable code
+    # which can be used in CI to test the docs without making network requests.
+    markdown = _add_path_to_code_blocks(markdown, page)
+
     if remove_base64_images:
         # Remove base64 encoded images from markdown
         markdown = re.sub(r"!\[.*?\]\(data:image/[^;]+;base64,[^\)]+\)", "", markdown)
@@ -160,6 +279,7 @@ def on_page_markdown(markdown: str, page: Page, **kwargs: Dict[str, Any]):
         **kwargs,
     )
 
+
 # redirects
 
 HTML_TEMPLATE = """

docs/_scripts/setup_vcr.py

@@ -0,0 +1,73 @@
+# A list of patterns that, if found in a code block, will cause us to leave that block unchanged.
+import hashlib
+import os
+from textwrap import dedent
+
+preambles = {
+    "python": "vcr_setup_preamble.py",
+    "typescript": "nock_setup_preamble.ts",
+}
+
+
+def _get_python_cassette_init(cassette_name: str, hash_: str) -> str:
+    return dedent(
+        f"""
+        _cassette = HashedCassette('{cassette_name}', '{hash_}')
+        _cassette.__enter__()
+        """
+    )
+
+
+def _get_typescript_cassette_init(cassette_name: str, hash_: str) -> str:
+    return dedent(
+        f"""
+        const _cassette = new HashedCassette("{cassette_name}", "{hash_}");
+        await _cassette.enter();
+        """
+    )
+
+
+def _get_python_cassette_cleanup() -> str:
+    return "_cassette.__exit__()"
+
+
+def _get_typescript_cassette_cleanup() -> str:
+    return "await _cassette.exit();"
+
+
+preamble_inits = {
+    "python": _get_python_cassette_init,
+    "typescript": _get_typescript_cassette_init,
+}
+
+preamble_cleanups = {
+    "python": _get_python_cassette_cleanup,
+    "typescript": _get_typescript_cassette_cleanup,
+}
+
+
+def load_preamble(language: str, code: str, cassette_name: str) -> str:
+    """Load the source code for the preamble for a given language."""
+    _assets_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "assets")
+
+    preamble_path = os.path.join(_assets_dir, preambles[language])
+    with open(preamble_path, "r") as f:
+        lines = f.readlines()
+        hash_ = _hash_string(code)
+        lines.append(preamble_inits[language](cassette_name, hash_))
+        return "\n".join(lines).strip()
+
+
+def load_postamble(language: str) -> str:
+    """Load the source code for the postamble for a given language."""
+
+    return preamble_cleanups[language]()
+
+
+def _hash_string(input_string: str) -> str:
+    # Encode the input string to bytes
+    encoded_string = input_string.encode("utf-8")
+    # Create a SHA-256 hash object
+    sha256_hash = hashlib.sha256(encoded_string)
+    # Get the hexadecimal digest of the hash
+    return sha256_hash.hexdigest()