mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-17 21:25:46 +02:00
238 lines
8.9 KiB
Python
238 lines
8.9 KiB
Python
"""Preprocess notebooks for CI. Currently adds VCR cassettes and optionally removes pip install cells."""
|
|
|
|
import logging
|
|
import os
|
|
import json
|
|
import click
|
|
import nbformat
|
|
|
|
logger = logging.getLogger(__name__)
|
|
NOTEBOOK_DIRS = ("docs/how-tos","docs/tutorials")
|
|
DOCS_PATH = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
|
CASSETTES_PATH = os.path.join(DOCS_PATH, "cassettes")
|
|
|
|
BLOCKLIST_COMMANDS = (
|
|
# skip if has WebBaseLoader to avoid caching web pages
|
|
"WebBaseLoader",
|
|
# skip if has draw_mermaid_png to avoid generating mermaid images via API
|
|
"draw_mermaid_png",
|
|
)
|
|
|
|
NOTEBOOKS_NO_CASSETTES = (
|
|
"docs/how-tos/visualization.ipynb",
|
|
"docs/how-tos/many-tools.ipynb"
|
|
)
|
|
|
|
NOTEBOOKS_NO_EXECUTION = [
|
|
# this uses a user provided project name for langsmith
|
|
"docs/tutorials/tnt-llm/tnt-llm.ipynb",
|
|
# this uses langsmith datasets
|
|
"docs/tutorials/chatbot-simulation-evaluation/langsmith-agent-simulation-evaluation.ipynb",
|
|
# this uses browser APIs
|
|
"docs/tutorials/web-navigation/web_voyager.ipynb",
|
|
# these RAG guides use an ollama model
|
|
"docs/tutorials/rag/langgraph_adaptive_rag_local.ipynb",
|
|
"docs/tutorials/rag/langgraph_crag_local.ipynb",
|
|
"docs/tutorials/rag/langgraph_self_rag_local.ipynb",
|
|
# this loads a massive dataset from gcp
|
|
"docs/tutorials/usaco/usaco.ipynb",
|
|
# TODO: figure out why autogen notebook is not runnable (they are just hanging. possible due to code execution?)
|
|
"docs/how-tos/autogen-integration.ipynb",
|
|
"docs/how-tos/autogen-integration-functional.ipynb",
|
|
# TODO: need to update these notebooks to make sure they are runnable in CI
|
|
"docs/tutorials/storm/storm.ipynb", # issues only when running with VCR
|
|
"docs/tutorials/lats/lats.ipynb", # issues only when running with VCR
|
|
"docs/tutorials/rag/langgraph_crag.ipynb", # flakiness from tavily
|
|
"docs/tutorials/rag/langgraph_adaptive_rag.ipynb", # flakiness only when running in GHA
|
|
"docs/tutorials/rag/langgraph_self_rag.ipynb", # flakiness only when running in GHA
|
|
"docs/tutorials/rag/langgraph_agentic_rag.ipynb", # flakiness only when running in GHA
|
|
"docs/how-tos/map-reduce.ipynb", # flakiness from structured output, only when running with VCR
|
|
"docs/tutorials/tot/tot.ipynb",
|
|
"docs/how-tos/visualization.ipynb",
|
|
"docs/tutorials/llm-compiler/LLMCompiler.ipynb"
|
|
]
|
|
|
|
|
|
def comment_install_cells(notebook: nbformat.NotebookNode) -> nbformat.NotebookNode:
|
|
for cell in notebook.cells:
|
|
if cell.cell_type != "code":
|
|
continue
|
|
|
|
if "pip install" in cell.source:
|
|
# Comment out the lines in cells containing "pip install"
|
|
cell.source = "\n".join(
|
|
f"# {line}" if line.strip() else line
|
|
for line in cell.source.splitlines()
|
|
)
|
|
|
|
return notebook
|
|
|
|
|
|
def is_magic_command(code: str) -> bool:
|
|
return code.strip().startswith("%") or code.strip().startswith("!")
|
|
|
|
|
|
def is_comment(code: str) -> bool:
|
|
return code.strip().startswith("#")
|
|
|
|
|
|
def has_blocklisted_command(code: str, metadata: dict) -> bool:
|
|
if 'hide_from_vcr' in metadata:
|
|
return True
|
|
|
|
code = code.strip()
|
|
for blocklisted_pattern in BLOCKLIST_COMMANDS:
|
|
if blocklisted_pattern in code:
|
|
return True
|
|
return False
|
|
|
|
|
|
def add_vcr_to_notebook(
|
|
notebook: nbformat.NotebookNode, cassette_prefix: str
|
|
) -> nbformat.NotebookNode:
|
|
"""Inject `with vcr.cassette` into each code cell of the notebook."""
|
|
|
|
uses_langsmith = False
|
|
# Inject VCR context manager into each code cell
|
|
for idx, cell in enumerate(notebook.cells):
|
|
if cell.cell_type != "code":
|
|
continue
|
|
|
|
lines = cell.source.splitlines()
|
|
# skip if empty cell
|
|
if not lines:
|
|
continue
|
|
|
|
are_magic_lines = [is_magic_command(line) for line in lines]
|
|
|
|
# skip if all magic
|
|
if all(are_magic_lines):
|
|
continue
|
|
|
|
if any(are_magic_lines):
|
|
raise ValueError(
|
|
"Cannot process code cells with mixed magic and non-magic code."
|
|
)
|
|
|
|
# skip if just comments
|
|
if all(is_comment(line) or not line.strip() for line in lines):
|
|
continue
|
|
|
|
if has_blocklisted_command(cell.source, cell.metadata):
|
|
continue
|
|
|
|
cell_id = cell.get("id", idx)
|
|
cassette_name = f"{cassette_prefix}_{cell_id}.msgpack.zlib"
|
|
cell.source = f"with custom_vcr.use_cassette('{cassette_name}', filter_headers=['x-api-key', 'authorization'], record_mode='once', serializer='advanced_compressed'):\n" + "\n".join(
|
|
f" {line}" for line in lines
|
|
)
|
|
|
|
if any("hub.pull" in line or "from langsmith import" in line for line in lines):
|
|
uses_langsmith = True
|
|
|
|
# Add import statement
|
|
vcr_import_lines = []
|
|
if uses_langsmith:
|
|
vcr_import_lines.extend([
|
|
# patch urllib3 to handle vcr errors, see more here:
|
|
# https://github.com/langchain-ai/langsmith-sdk/blob/main/python/langsmith/_internal/_patch.py
|
|
"import sys",
|
|
f"sys.path.insert(0, '{os.path.join(DOCS_PATH, '_scripts')}')",
|
|
"import _patch as patch_urllib3",
|
|
"patch_urllib3.patch_urllib3()",
|
|
])
|
|
|
|
vcr_import_lines.extend([
|
|
"import nest_asyncio",
|
|
"nest_asyncio.apply()",
|
|
"import vcr",
|
|
"import msgpack",
|
|
"import base64",
|
|
"import zlib",
|
|
"import os",
|
|
"os.environ.pop(\"LANGCHAIN_TRACING_V2\", None)",
|
|
"custom_vcr = vcr.VCR()",
|
|
"",
|
|
"def compress_data(data, compression_level=9):",
|
|
" 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):",
|
|
" decoded = base64.b64decode(compressed_string)",
|
|
" decompressed = zlib.decompress(decoded)",
|
|
" return msgpack.unpackb(decompressed, raw=False)",
|
|
"",
|
|
"class AdvancedCompressedSerializer:",
|
|
" def serialize(self, cassette_dict):",
|
|
" return compress_data(cassette_dict)",
|
|
"",
|
|
" def deserialize(self, cassette_string):",
|
|
" return decompress_data(cassette_string)",
|
|
"",
|
|
"custom_vcr.register_serializer('advanced_compressed', AdvancedCompressedSerializer())",
|
|
"custom_vcr.serializer = 'advanced_compressed'",
|
|
])
|
|
|
|
import_cell = nbformat.v4.new_code_cell(source="\n".join(vcr_import_lines))
|
|
import_cell.pop("id", None)
|
|
notebook.cells.insert(0, import_cell)
|
|
return notebook
|
|
|
|
|
|
def process_notebooks(should_comment_install_cells: bool) -> None:
|
|
for directory in NOTEBOOK_DIRS:
|
|
for root, _, files in os.walk(directory):
|
|
for file in files:
|
|
if not file.endswith(".ipynb") or "ipynb_checkpoints" in root:
|
|
continue
|
|
|
|
notebook_path = os.path.join(root, file)
|
|
try:
|
|
notebook = nbformat.read(notebook_path, as_version=4)
|
|
|
|
if should_comment_install_cells:
|
|
notebook = comment_install_cells(notebook)
|
|
|
|
base_filename = os.path.splitext(os.path.basename(file))[0]
|
|
cassette_prefix = os.path.join(CASSETTES_PATH, base_filename)
|
|
if notebook_path not in NOTEBOOKS_NO_CASSETTES:
|
|
notebook = add_vcr_to_notebook(
|
|
notebook, cassette_prefix=cassette_prefix
|
|
)
|
|
|
|
if notebook_path in NOTEBOOKS_NO_EXECUTION:
|
|
# Add a cell at the beginning to indicate that this notebook should not be executed
|
|
warning_cell = nbformat.v4.new_markdown_cell(
|
|
source="**Warning:** This notebook is not meant to be executed automatically."
|
|
)
|
|
notebook.cells.insert(0, warning_cell)
|
|
|
|
# Add a special tag to the first code cell
|
|
if notebook.cells and notebook.cells[1].cell_type == "code":
|
|
notebook.cells[1].metadata["tags"] = notebook.cells[1].metadata.get("tags", []) + ["no_execution"]
|
|
|
|
nbformat.write(notebook, notebook_path)
|
|
logger.info(f"Processed: {notebook_path}")
|
|
except Exception as e:
|
|
logger.error(f"Error processing {notebook_path}: {e}")
|
|
|
|
with open("notebooks_no_execution.json", "w") as f:
|
|
json.dump(NOTEBOOKS_NO_EXECUTION, f)
|
|
|
|
|
|
@click.command()
|
|
@click.option(
|
|
"--comment-install-cells",
|
|
is_flag=True,
|
|
default=False,
|
|
help="Whether to comment out install cells",
|
|
)
|
|
def main(comment_install_cells):
|
|
process_notebooks(should_comment_install_cells=comment_install_cells)
|
|
logger.info("All notebooks processed successfully.")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|