diff --git a/docs/_scripts/generate_api_reference_links.py b/docs/_scripts/generate_api_reference_links.py index 9ba0978bf..dd8f0b662 100644 --- a/docs/_scripts/generate_api_reference_links.py +++ b/docs/_scripts/generate_api_reference_links.py @@ -1,9 +1,9 @@ +import ast import importlib -import inspect import logging import re from functools import lru_cache -from typing import List, Literal, Optional +from typing import List, Optional from typing_extensions import TypedDict @@ -39,7 +39,6 @@ MANUAL_API_REFERENCES_LANGGRAPH = [ (["langgraph.graph"], "langgraph.graph.message", "add_messages", "graphs"), (["langgraph.graph"], "langgraph.graph.state", "StateGraph", "graphs"), (["langgraph.graph"], "langgraph.graph.state", "CompiledStateGraph", "graphs"), - ([], "langgraph.types", "StreamMode", "types"), (["langgraph.graph"], "langgraph.constants", "START", "constants"), (["langgraph.graph"], "langgraph.constants", "END", "constants"), (["langgraph.constants"], "langgraph.types", "Send", "types"), @@ -48,7 +47,9 @@ MANUAL_API_REFERENCES_LANGGRAPH = [ (["langgraph.constants"], "langgraph.types", "Command", "types"), (["langgraph.func"], "langgraph.func", "entrypoint", "func"), (["langgraph.func"], "langgraph.func", "task", "func"), - ([], "langgraph.types", "RetryPolicy", "types"), + (["langgraph.types"], "langgraph.types", "RetryPolicy", "types"), + (["langgraph.types"], "langgraph.types", "StreamMode", "types"), + (["langgraph.types"], "langgraph.types", "StreamWriter", "types"), ([], "langgraph.checkpoint.base", "Checkpoint", "checkpoints"), ([], "langgraph.checkpoint.base", "CheckpointMetadata", "checkpoints"), ([], "langgraph.checkpoint.base", "BaseCheckpointSaver", "checkpoints"), @@ -68,34 +69,19 @@ WELL_KNOWN_LANGGRAPH_OBJECTS = { } -def _make_regular_expression(pkg_prefix: str) -> re.Pattern: - if not pkg_prefix.isidentifier(): - raise ValueError(f"Invalid package prefix: {pkg_prefix}") - return re.compile( - r"from\s+(" + pkg_prefix + "(?:_\w+)?(?:\.\w+)*?)\s+import\s+" - r"((?:\w+(?:,\s*)?)*" # Match zero or more words separated by a comma+optional ws - r"(?:\s*\(.*?\))?)", # Match optional parentheses block - re.DOTALL, # Match newlines as well - ) - - -# Regular expression to match langchain import lines -_IMPORT_LANGCHAIN_RE = _make_regular_expression("langchain") -_IMPORT_LANGGRAPH_RE = _make_regular_expression("langgraph") - - @lru_cache(maxsize=10_000) def _get_full_module_name(module_path: str, class_name: str) -> Optional[str]: """Get full module name using inspect, with LRU cache to memoize results.""" try: module = importlib.import_module(module_path) - class_ = getattr(module, class_name) - module = inspect.getmodule(class_) - if module is None: - # For constants, inspect.getmodule() might return None - # In this case, we'll return the original module_path + symbol = getattr(module, class_name) + # First check the __module__ attribute on the symbol. + mod_name = getattr(symbol, "__module__", None) + # If __module__ is not set or comes from typing, + # assume the definition is in module_path. + if mod_name is None or mod_name.startswith("typing"): return module_path - return module.__name__ + return mod_name except AttributeError as e: logger.warning(f"API Reference: Could not find module for {class_name}, {e}") return None @@ -104,139 +90,128 @@ def _get_full_module_name(module_path: str, class_name: str) -> Optional[str]: return None -def _get_doc_title(data: str, file_name: str) -> str: - try: - return re.findall(r"^#\s*(.*)", data, re.MULTILINE)[0] - except IndexError: - pass - # Parse the rst-style titles - try: - return re.findall(r"^(.*)\n=+\n", data, re.MULTILINE)[0] - except IndexError: - return file_name - - class ImportInformation(TypedDict): imported: str # The name of the class that was imported. source: str # The full module path from which the class was imported. docs: str # The URL pointing to the class's documentation. - title: str # The title of the document where the import is used. + path: str # The path of the file where the markdown content originated. -def _get_imports( - code: str, doc_title: str, package_ecosystem: Literal["langchain", "langgraph"] -) -> List[ImportInformation]: - """Get imports from the given code block. - - Args: - code: Python code block from which to extract imports - doc_title: Title of the document - package_ecosystem: "langchain" or "langgraph". The two live in different - repositories and have separate documentation sites. - - Returns: - List of import information for the given code block - """ - imports = [] - - if package_ecosystem == "langchain": - pattern = _IMPORT_LANGCHAIN_RE - elif package_ecosystem == "langgraph": - pattern = _IMPORT_LANGGRAPH_RE - else: - raise ValueError(f"Invalid package ecosystem: {package_ecosystem}") - - for import_match in pattern.finditer(code): - module = import_match.group(1) - if "pydantic_v1" in module: - continue - imports_str = ( - import_match.group(2).replace("(\n", "").replace("\n)", "") - ) # Handle newlines within parentheses - # remove any newline and spaces, then split by comma - imported_classes = [ - imp.strip() - for imp in re.split(r",\s*", imports_str.replace("\n", "")) - if imp.strip() - ] - for class_name in imported_classes: - module_path = _get_full_module_name(module, class_name) - if not module_path: - continue - if len(module_path.split(".")) < 2: - continue - - if package_ecosystem == "langchain": - pkg = module_path.split(".")[0].replace("langchain_", "") - top_level_mod = module_path.split(".")[1] - - url = ( - _LANGCHAIN_API_REFERENCE - + pkg - + "/" - + top_level_mod - + "/" - + module_path - + "." - + class_name - + ".html" - ) - elif package_ecosystem == "langgraph": - if (module, class_name) not in WELL_KNOWN_LANGGRAPH_OBJECTS: - # Likely not documented yet - continue - - source_module, namespace = WELL_KNOWN_LANGGRAPH_OBJECTS[ - (module, class_name) - ] - url = ( - _LANGGRAPH_API_REFERENCE - + namespace - + "/#" - + source_module - + "." - + class_name - ) - else: - raise ValueError(f"Invalid package ecosystem: {package_ecosystem}") - - # Add the import information to our list - imports.append( - { - "imported": class_name, - "source": module, - "docs": url, - "title": doc_title, - } - ) - - return imports - - -def get_imports(code: str, doc_title: str) -> List[ImportInformation]: +def get_imports(code: str, path: str) -> List[ImportInformation]: """Retrieve all import references from the given code for specified ecosystems. Args: code: The source code from which to extract import references. - doc_title: The documentation title associated with the code. + path: The path of the file where the markdown content originated. Returns: A list of import information for each import found. """ - ecosystems = ["langchain", "langgraph"] - all_imports = [] - for package_ecosystem in ecosystems: - all_imports.extend(_get_imports(code, doc_title, package_ecosystem)) - return all_imports + # Parse the code into an AST. + try: + tree = ast.parse(code) + except SyntaxError: + return [] + + found_imports = [] + + # Walk through the AST and process ImportFrom nodes. + for node in ast.walk(tree): + if isinstance(node, ast.ImportFrom): + # node.module is the source module. + if node.module is None: + continue + for alias in node.names: + if not ( + node.module.startswith("langchain") + or node.module.startswith("langgraph") + ): + continue + + found_imports.append( + { + "source": node.module, + # alias.name is the original name even if an alias exists. + "imported": alias.name, + } + ) + + imports: list[ImportInformation] = [] + + for found_import in found_imports: + module = found_import["source"] + + if module.startswith("langchain"): + # Handles things like `langchain` or `langchain_anthropic` + package_ecosystem = "langchain" + elif module.startswith("langgraph"): + package_ecosystem = "langgraph" + else: + continue + + class_name = found_import["imported"] + module_path = _get_full_module_name(module, class_name) + if not module_path: + continue + if len(module_path.split(".")) < 2: + continue + + if package_ecosystem == "langchain": + pkg = module_path.split(".")[0].replace("langchain_", "") + top_level_mod = module_path.split(".")[1] + + url = ( + _LANGCHAIN_API_REFERENCE + + pkg + + "/" + + top_level_mod + + "/" + + module_path + + "." + + class_name + + ".html" + ) + elif package_ecosystem == "langgraph": + if (module, class_name) not in WELL_KNOWN_LANGGRAPH_OBJECTS: + # Likely not documented yet + continue + + source_module, namespace = WELL_KNOWN_LANGGRAPH_OBJECTS[ + (module, class_name) + ] + url = ( + _LANGGRAPH_API_REFERENCE + + namespace + + "/#" + + source_module + + "." + + class_name + ) + else: + raise ValueError(f"Invalid package ecosystem: {package_ecosystem}") + + # Add the import information to our list + imports.append( + { + "imported": class_name, + "source": module, + "docs": url, + "path": path, + } + ) + + return imports -def update_markdown_with_imports(markdown: str) -> str: +def update_markdown_with_imports(markdown: str, path: str) -> str: """Update markdown to include API reference links for imports in Python code blocks. - This function scans the markdown content for Python code blocks, extracts any imports, and appends links to their API documentation. + This function scans the markdown content for Python code blocks, extracts any + imports, and appends links to their API documentation. Args: markdown: The markdown content to process. + path: The path of the file where the markdown content originated. Returns: Updated markdown with API reference links appended to Python code blocks. @@ -247,7 +222,8 @@ def update_markdown_with_imports(markdown: str) -> str: ```python from langchain.nlp import TextGenerator ``` - This function will append an API reference link to the `TextGenerator` class from the `langchain.nlp` module if it's recognized. + This function will append an API reference link to the `TextGenerator` class + from the `langchain.nlp` module if it's recognized. """ code_block_pattern = re.compile( r"(?P[ \t]*)```(?Ppython|py)\n(?P.*?)\n(?P=indent)```", diff --git a/docs/_scripts/notebook_hooks.py b/docs/_scripts/notebook_hooks.py index 5089e50d6..bc6ec5438 100644 --- a/docs/_scripts/notebook_hooks.py +++ b/docs/_scripts/notebook_hooks.py @@ -175,7 +175,7 @@ def _on_page_markdown_with_config( # Append API reference links to code blocks if add_api_references: - markdown = update_markdown_with_imports(markdown) + markdown = update_markdown_with_imports(markdown, page.file.abs_src_path) # Apply highlight comments to code blocks markdown = _highlight_code_blocks(markdown) diff --git a/docs/tests/unit_tests/test_api_reference.py b/docs/tests/unit_tests/test_api_reference.py new file mode 100644 index 000000000..0d3bc57a9 --- /dev/null +++ b/docs/tests/unit_tests/test_api_reference.py @@ -0,0 +1,212 @@ +"""Test generation of links into the API reference.""" + +import pytest + +from _scripts.generate_api_reference_links import ( + update_markdown_with_imports, + get_imports, +) + +MARKDOWN_IMPORTS = """\ +```python +from langgraph.types import interrupt +``` +""" + +EXPECTED_MARKDOWN = """\ +```python +from langgraph.types import interrupt +``` + +API Reference: interrupt +""" + + +def test_update_markdown_with_imports() -> None: + """Light weight end-to-end test.""" + assert ( + update_markdown_with_imports(MARKDOWN_IMPORTS, "some_path") == EXPECTED_MARKDOWN + ) + + +@pytest.mark.parametrize( + "code_block, expected_imports", + [ + ( + "from langgraph.types import interrupt", + [ + { + "docs": "https://langchain-ai.github.io/langgraph/reference/types/#langgraph.types.interrupt", + "imported": "interrupt", + "path": "some_path", + "source": "langgraph.types", + } + ], + ), + ( + "from langgraph.types import ( interrupt )", + [ + { + "docs": "https://langchain-ai.github.io/langgraph/reference/types/#langgraph.types.interrupt", + "imported": "interrupt", + "path": "some_path", + "source": "langgraph.types", + } + ], + ), + ( + "from langgraph.types import interrupt as foo", + [ + { + "docs": "https://langchain-ai.github.io/langgraph/reference/types/#langgraph.types.interrupt", + "imported": "interrupt", + "path": "some_path", + "source": "langgraph.types", + } + ], + ), + ], +) +def test_get_imports(code_block: str, expected_imports: list) -> None: + """Get imports from a code block.""" + assert ( + get_imports(code_block, "some_path") == expected_imports + ), f"Failed for code_block=`{code_block}`" + + +@pytest.mark.parametrize( + "code, expected_imports", + [ + # Single import without parenthesis + ( + "from langgraph.types import interrupt", + [ + { + "source": "langgraph.types", + "imported": "interrupt", + } + ], + ), + # Multiple imports + ( + ( + "from langgraph.types import interrupt\n" + "from langgraph.func import task" + ), + [ + { + "source": "langgraph.types", + "imported": "interrupt", + }, + { + "source": "langgraph.func", + "imported": "task", + }, + ], + ), + # Single import with parenthesis and extra whitespace + ( + "from langgraph.types import ( interrupt )", + [ + { + "source": "langgraph.types", + "imported": "interrupt", + } + ], + ), + # Single import with an alias + ( + "from langgraph.types import interrupt as foo", + [ + { + "source": "langgraph.types", + "imported": "interrupt", + } + ], + ), + # Multiple imports on one line with an alias + ( + "from langgraph.types import interrupt, StreamWriter as bar", + [ + { + "source": "langgraph.types", + "imported": "interrupt", + }, + { + "source": "langgraph.types", + "imported": "StreamWriter", + }, + ], + ), + # Multiple imports without aliases + ( + "from langgraph.types import interrupt, StreamWriter", + [ + { + "source": "langgraph.types", + "imported": "interrupt", + }, + { + "source": "langgraph.types", + "imported": "StreamWriter", + }, + ], + ), + # Multiline import with parenthesis and trailing comma + ( + """from langgraph.types import ( + interrupt, + StreamWriter as foo, + Command, + )""", + [ + { + "source": "langgraph.types", + "imported": "interrupt", + }, + { + "source": "langgraph.types", + "imported": "StreamWriter", + }, + { + "source": "langgraph.types", + "imported": "Command", + }, + ], + ), + # Multiline import with parenthesis and trailing comma + ( + ( + "from langgraph.types import (\n" + " interrupt,\n" + " StreamWriter as foo\n," + " Command,\n" + ")\n" + "def foo():\n" + " pass\n" + "" + ), + [ + { + "source": "langgraph.types", + "imported": "interrupt", + }, + { + "source": "langgraph.types", + "imported": "StreamWriter", + }, + { + "source": "langgraph.types", + "imported": "Command", + }, + ], + ), + ], +) +def test_regexp_matching(code: str, expected_imports: list) -> None: + results = get_imports(code, "some_path") + for result in results: + del result["docs"] + del result["path"] + + assert results == expected_imports