langchain-core 0.4.0.dev0__py3-none-any.whl → 1.0.0__py3-none-any.whl

langchain_core/runnables/graph_ascii.py

@@ -3,12 +3,24 @@
 Adapted from https://github.com/iterative/dvc/blob/main/dvc/dagascii.py.
 """
 
+from __future__ import annotations
+
 import math
 import os
 from collections.abc import Mapping, Sequence
-from typing import Any
+from typing import TYPE_CHECKING, Any
+
+try:
+    from grandalf.graphs import Edge, Graph, Vertex  # type: ignore[import-untyped]
+    from grandalf.layouts import SugiyamaLayout  # type: ignore[import-untyped]
+    from grandalf.routing import route_with_lines  # type: ignore[import-untyped]
+
+    _HAS_GRANDALF = True
+except ImportError:
+    _HAS_GRANDALF = False
 
-from langchain_core.runnables.graph import Edge as LangEdge
+if TYPE_CHECKING:
+    from langchain_core.runnables.graph import Edge as LangEdge
 
 
 class VertexViewer:
@@ -50,8 +62,11 @@ class AsciiCanvas:
         """Create an ASCII canvas.
 
         Args:
-            cols (int): number of columns in the canvas. Should be > 1.
-            lines (int): number of lines in the canvas. Should be > 1.
+            cols: number of columns in the canvas. Should be `> 1`.
+            lines: number of lines in the canvas. Should be `> 1`.
+
+        Raises:
+            ValueError: if canvas dimensions are invalid.
         """
         if cols <= 1 or lines <= 1:
             msg = "Canvas dimensions should be > 1"
@@ -63,7 +78,11 @@ class AsciiCanvas:
         self.canvas = [[" "] * cols for line in range(lines)]
 
     def draw(self) -> str:
-        """Draws ASCII canvas on the screen."""
+        """Draws ASCII canvas on the screen.
+
+        Returns:
+            The ASCII canvas string.
+        """
         lines = map("".join, self.canvas)
         return os.linesep.join(lines)
 
@@ -71,12 +90,16 @@ class AsciiCanvas:
         """Create a point on ASCII canvas.
 
         Args:
-            x (int): x coordinate. Should be >= 0 and < number of columns in
+            x: x coordinate. Should be `>= 0` and `<` number of columns in
                 the canvas.
-            y (int): y coordinate. Should be >= 0 an < number of lines in the
+            y: y coordinate. Should be `>= 0` an `<` number of lines in the
                 canvas.
-            char (str): character to place in the specified point on the
+            char: character to place in the specified point on the
                 canvas.
+
+        Raises:
+            ValueError: if char is not a single character or if
+                coordinates are out of bounds.
         """
         if len(char) != 1:
             msg = "char should be a single character"
@@ -94,11 +117,11 @@ class AsciiCanvas:
         """Create a line on ASCII canvas.
 
         Args:
-            x0 (int): x coordinate where the line should start.
-            y0 (int): y coordinate where the line should start.
-            x1 (int): x coordinate where the line should end.
-            y1 (int): y coordinate where the line should end.
-            char (str): character to draw the line with.
+            x0: x coordinate where the line should start.
+            y0: y coordinate where the line should start.
+            x1: x coordinate where the line should end.
+            y1: y coordinate where the line should end.
+            char: character to draw the line with.
         """
         if x0 > x1:
             x1, x0 = x0, x1
@@ -126,9 +149,9 @@ class AsciiCanvas:
         """Print a text on ASCII canvas.
 
         Args:
-            x (int): x coordinate where the text should start.
-            y (int): y coordinate where the text should start.
-            text (str): string that should be printed.
+            x: x coordinate where the text should start.
+            y: y coordinate where the text should start.
+            text: string that should be printed.
         """
         for i, char in enumerate(text):
             self.point(x + i, y, char)
@@ -137,10 +160,10 @@ class AsciiCanvas:
         """Create a box on ASCII canvas.
 
         Args:
-            x0 (int): x coordinate of the box corner.
-            y0 (int): y coordinate of the box corner.
-            width (int): box width.
-            height (int): box height.
+            x0: x coordinate of the box corner.
+            y0: y coordinate of the box corner.
+            width: box width.
+            height: box height.
         """
         if width <= 1 or height <= 1:
             msg = "Box dimensions should be > 1"
@@ -174,13 +197,9 @@ class _EdgeViewer:
 def _build_sugiyama_layout(
     vertices: Mapping[str, str], edges: Sequence[LangEdge]
 ) -> Any:
-    try:
-        from grandalf.graphs import Edge, Graph, Vertex  # type: ignore[import-untyped]
-        from grandalf.layouts import SugiyamaLayout  # type: ignore[import-untyped]
-        from grandalf.routing import route_with_lines  # type: ignore[import-untyped]
-    except ImportError as exc:
+    if not _HAS_GRANDALF:
         msg = "Install grandalf to draw graphs: `pip install grandalf`."
-        raise ImportError(msg) from exc
+        raise ImportError(msg)
 
     #
     # Just a reminder about naming conventions:
@@ -225,28 +244,31 @@ def draw_ascii(vertices: Mapping[str, str], edges: Sequence[LangEdge]) -> str:
     """Build a DAG and draw it in ASCII.
 
     Args:
-        vertices (list): list of graph vertices.
-        edges (list): list of graph edges.
+        vertices: list of graph vertices.
+        edges: list of graph edges.
+
+    Raises:
+        ValueError: if the canvas dimensions are invalid or if
+            edge coordinates are invalid.
 
     Returns:
-        str: ASCII representation
+        ASCII representation
 
     Example:
+        ```python
+        from langchain_core.runnables.graph_ascii import draw_ascii
 
-        .. code-block:: python
-
-            from langchain_core.runnables.graph_ascii import draw_ascii
+        vertices = {1: "1", 2: "2", 3: "3", 4: "4"}
+        edges = [
+            (source, target, None, None)
+            for source, target in [(1, 2), (2, 3), (2, 4), (1, 4)]
+        ]
 
-            vertices = {1: "1", 2: "2", 3: "3", 4: "4"}
-            edges = [
-                (source, target, None, None)
-                for source, target in [(1, 2), (2, 3), (2, 4), (1, 4)]
-            ]
 
+        print(draw_ascii(vertices, edges))
+        ```
 
-            print(draw_ascii(vertices, edges))
-
-        .. code-block:: none
+        ```txt
 
                  +---+
                  | 1 |
@@ -263,7 +285,7 @@ def draw_ascii(vertices: Mapping[str, str], edges: Sequence[LangEdge]) -> str:
             +---+     +---+
             | 3 |     | 4 |
             +---+     +---+
-
+        ```
     """
     # NOTE: coordinates might me negative, so we need to shift
     # everything to the positive plane before we actually draw it.

langchain_core/runnables/graph_mermaid.py

@@ -1,24 +1,43 @@
 """Mermaid graph drawing utilities."""
 
+from __future__ import annotations
+
 import asyncio
 import base64
 import random
 import re
+import string
 import time
 from dataclasses import asdict
 from pathlib import Path
-from typing import Any, Literal, Optional
+from typing import TYPE_CHECKING, Any, Literal
 
 import yaml
 
 from langchain_core.runnables.graph import (
     CurveStyle,
-    Edge,
     MermaidDrawMethod,
-    Node,
     NodeStyles,
 )
 
+if TYPE_CHECKING:
+    from langchain_core.runnables.graph import Edge, Node
+
+
+try:
+    import requests
+
+    _HAS_REQUESTS = True
+except ImportError:
+    _HAS_REQUESTS = False
+
+try:
+    from pyppeteer import launch  # type: ignore[import-not-found]
+
+    _HAS_PYPPETEER = True
+except ImportError:
+    _HAS_PYPPETEER = False
+
 MARKDOWN_SPECIAL_CHARS = "*_`"
 
 
@@ -26,50 +45,44 @@ def draw_mermaid(
     nodes: dict[str, Node],
     edges: list[Edge],
     *,
-    first_node: Optional[str] = None,
-    last_node: Optional[str] = None,
+    first_node: str | None = None,
+    last_node: str | None = None,
     with_styles: bool = True,
     curve_style: CurveStyle = CurveStyle.LINEAR,
-    node_styles: Optional[NodeStyles] = None,
+    node_styles: NodeStyles | None = None,
     wrap_label_n_words: int = 9,
-    frontmatter_config: Optional[dict[str, Any]] = None,
+    frontmatter_config: dict[str, Any] | None = None,
 ) -> str:
     """Draws a Mermaid graph using the provided graph data.
 
     Args:
-        nodes (dict[str, str]): List of node ids.
-        edges (list[Edge]): List of edges, object with a source,
-            target and data.
-        first_node (str, optional): Id of the first node. Defaults to None.
-        last_node (str, optional): Id of the last node. Defaults to None.
-        with_styles (bool, optional): Whether to include styles in the graph.
-            Defaults to True.
-        curve_style (CurveStyle, optional): Curve style for the edges.
-            Defaults to CurveStyle.LINEAR.
-        node_styles (NodeStyles, optional): Node colors for different types.
-            Defaults to NodeStyles().
-        wrap_label_n_words (int, optional): Words to wrap the edge labels.
-            Defaults to 9.
-        frontmatter_config (dict[str, Any], optional): Mermaid frontmatter config.
+        nodes: List of node ids.
+        edges: List of edges, object with a source, target and data.
+        first_node: Id of the first node.
+        last_node: Id of the last node.
+        with_styles: Whether to include styles in the graph.
+        curve_style: Curve style for the edges.
+        node_styles: Node colors for different types.
+        wrap_label_n_words: Words to wrap the edge labels.
+        frontmatter_config: Mermaid frontmatter config.
             Can be used to customize theme and styles. Will be converted to YAML and
-            added to the beginning of the mermaid graph. Defaults to None.
+            added to the beginning of the mermaid graph.
 
             See more here: https://mermaid.js.org/config/configuration.html.
 
             Example config:
 
-            .. code-block:: python
-
+            ```python
             {
                 "config": {
                     "theme": "neutral",
                     "look": "handDrawn",
-                    "themeVariables": { "primaryColor": "#e2e2e2"},
+                    "themeVariables": {"primaryColor": "#e2e2e2"},
                 }
             }
-
+            ```
     Returns:
-        str: Mermaid graph syntax.
+        Mermaid graph syntax.
 
     """
     # Initialize Mermaid graph configuration
@@ -130,7 +143,7 @@ def draw_mermaid(
                 + "</em></small>"
             )
         node_label = format_dict.get(key, format_dict[default_class_label]).format(
-            _escape_node_label(key), label
+            _to_safe_id(key), label
         )
         return f"{indent}{node_label}\n"
 
@@ -145,7 +158,7 @@ def draw_mermaid(
         src_parts = edge.source.split(":")
         tgt_parts = edge.target.split(":")
         common_prefix = ":".join(
-            src for src, tgt in zip(src_parts, tgt_parts) if src == tgt
+            src for src, tgt in zip(src_parts, tgt_parts, strict=False) if src == tgt
         )
         edge_groups.setdefault(common_prefix, []).append(edge)
 
@@ -155,7 +168,7 @@ def draw_mermaid(
         nonlocal mermaid_graph
         self_loop = len(edges) == 1 and edges[0].source == edges[0].target
         if prefix and not self_loop:
-            subgraph = prefix.split(":")[-1]
+            subgraph = prefix.rsplit(":", maxsplit=1)[-1]
             if subgraph in seen_subgraphs:
                 msg = (
                     f"Found duplicate subgraph '{subgraph}' -- this likely means that "
@@ -193,8 +206,7 @@ def draw_mermaid(
                 edge_label = " -.-> " if edge.conditional else " --> "
 
             mermaid_graph += (
-                f"\t{_escape_node_label(source)}{edge_label}"
-                f"{_escape_node_label(target)};\n"
+                f"\t{_to_safe_id(source)}{edge_label}{_to_safe_id(target)};\n"
             )
 
         # Recursively add nested subgraphs
@@ -214,7 +226,7 @@ def draw_mermaid(
 
     # Add remaining subgraphs with edges
     for prefix, edges_ in edge_groups.items():
-        if ":" in prefix or prefix == "":
+        if not prefix or ":" in prefix:
             continue
         add_subgraph(edges_, prefix)
         seen_subgraphs.add(prefix)
@@ -238,9 +250,18 @@ def draw_mermaid(
     return mermaid_graph
 
 
-def _escape_node_label(node_label: str) -> str:
-    """Escapes the node label for Mermaid syntax."""
-    return re.sub(r"[^a-zA-Z-_0-9]", "_", node_label)
+def _to_safe_id(label: str) -> str:
+    """Convert a string into a Mermaid-compatible node id.
+
+    Keep [a-zA-Z0-9_-] characters unchanged.
+    Map every other character -> backslash + lowercase hex codepoint.
+
+    Result is guaranteed to be unique and Mermaid-compatible,
+    so nodes with special characters always render correctly.
+    """
+    allowed = string.ascii_letters + string.digits + "_-"
+    out = [ch if ch in allowed else "\\" + format(ord(ch), "x") for ch in label]
+    return "".join(out)
 
 
 def _generate_mermaid_graph_styles(node_colors: NodeStyles) -> str:
@@ -253,38 +274,33 @@ def _generate_mermaid_graph_styles(node_colors: NodeStyles) -> str:
 
 def draw_mermaid_png(
     mermaid_syntax: str,
-    output_file_path: Optional[str] = None,
+    output_file_path: str | None = None,
     draw_method: MermaidDrawMethod = MermaidDrawMethod.API,
-    background_color: Optional[str] = "white",
+    background_color: str | None = "white",
     padding: int = 10,
     max_retries: int = 1,
     retry_delay: float = 1.0,
+    base_url: str | None = None,
 ) -> bytes:
     """Draws a Mermaid graph as PNG using provided syntax.
 
     Args:
-        mermaid_syntax (str): Mermaid graph syntax.
-        output_file_path (str, optional): Path to save the PNG image.
-            Defaults to None.
-        draw_method (MermaidDrawMethod, optional): Method to draw the graph.
-            Defaults to MermaidDrawMethod.API.
-        background_color (str, optional): Background color of the image.
-            Defaults to "white".
-        padding (int, optional): Padding around the image. Defaults to 10.
-        max_retries (int, optional): Maximum number of retries (MermaidDrawMethod.API).
-            Defaults to 1.
-        retry_delay (float, optional): Delay between retries (MermaidDrawMethod.API).
-            Defaults to 1.0.
+        mermaid_syntax: Mermaid graph syntax.
+        output_file_path: Path to save the PNG image.
+        draw_method: Method to draw the graph.
+        background_color: Background color of the image.
+        padding: Padding around the image.
+        max_retries: Maximum number of retries (MermaidDrawMethod.API).
+        retry_delay: Delay between retries (MermaidDrawMethod.API).
+        base_url: Base URL for the Mermaid.ink API.
 
     Returns:
-        bytes: PNG image bytes.
+        PNG image bytes.
 
     Raises:
         ValueError: If an invalid draw method is provided.
     """
     if draw_method == MermaidDrawMethod.PYPPETEER:
-        import asyncio
-
         img_bytes = asyncio.run(
             _render_mermaid_using_pyppeteer(
                 mermaid_syntax, output_file_path, background_color, padding
@@ -297,6 +313,7 @@ def draw_mermaid_png(
             background_color=background_color,
             max_retries=max_retries,
             retry_delay=retry_delay,
+            base_url=base_url,
         )
     else:
         supported_methods = ", ".join([m.value for m in MermaidDrawMethod])
@@ -311,17 +328,15 @@ def draw_mermaid_png(
 
 async def _render_mermaid_using_pyppeteer(
     mermaid_syntax: str,
-    output_file_path: Optional[str] = None,
-    background_color: Optional[str] = "white",
+    output_file_path: str | None = None,
+    background_color: str | None = "white",
     padding: int = 10,
     device_scale_factor: int = 3,
 ) -> bytes:
     """Renders Mermaid graph using Pyppeteer."""
-    try:
-        from pyppeteer import launch  # type: ignore[import-not-found]
-    except ImportError as e:
+    if not _HAS_PYPPETEER:
         msg = "Install Pyppeteer to use the Pyppeteer method: `pip install pyppeteer`."
-        raise ImportError(msg) from e
+        raise ImportError(msg)
 
     browser = await launch()
     page = await browser.newPage()
@@ -385,21 +400,23 @@ async def _render_mermaid_using_pyppeteer(
 def _render_mermaid_using_api(
     mermaid_syntax: str,
     *,
-    output_file_path: Optional[str] = None,
-    background_color: Optional[str] = "white",
-    file_type: Optional[Literal["jpeg", "png", "webp"]] = "png",
+    output_file_path: str | None = None,
+    background_color: str | None = "white",
+    file_type: Literal["jpeg", "png", "webp"] | None = "png",
     max_retries: int = 1,
     retry_delay: float = 1.0,
+    base_url: str | None = None,
 ) -> bytes:
     """Renders Mermaid graph using the Mermaid.INK API."""
-    try:
-        import requests
-    except ImportError as e:
+    # Defaults to using the public mermaid.ink server.
+    base_url = base_url if base_url is not None else "https://mermaid.ink"
+
+    if not _HAS_REQUESTS:
         msg = (
             "Install the `requests` module to use the Mermaid.INK API: "
             "`pip install requests`."
         )
-        raise ImportError(msg) from e
+        raise ImportError(msg)
 
     # Use Mermaid API to render the image
     mermaid_syntax_encoded = base64.b64encode(mermaid_syntax.encode("utf8")).decode(
@@ -413,7 +430,7 @@ def _render_mermaid_using_api(
             background_color = f"!{background_color}"
 
     image_url = (
-        f"https://mermaid.ink/img/{mermaid_syntax_encoded}"
+        f"{base_url}/img/{mermaid_syntax_encoded}"
         f"?type={file_type}&bgColor={background_color}"
     )
 
@@ -445,7 +462,7 @@ def _render_mermaid_using_api(
 
             # For other status codes, fail immediately
             msg = (
-                "Failed to reach https://mermaid.ink/ API while trying to render "
+                f"Failed to reach {base_url} API while trying to render "
                 f"your graph. Status code: {response.status_code}.\n\n"
             ) + error_msg_suffix
             raise ValueError(msg)
@@ -457,14 +474,14 @@ def _render_mermaid_using_api(
                 time.sleep(sleep_time)
             else:
                 msg = (
-                    "Failed to reach https://mermaid.ink/ API while trying to render "
+                    f"Failed to reach {base_url} API while trying to render "
                     f"your graph after {max_retries} retries. "
                 ) + error_msg_suffix
                 raise ValueError(msg) from e
 
     # This should not be reached, but just in case
     msg = (
-        "Failed to reach https://mermaid.ink/ API while trying to render "
+        f"Failed to reach {base_url} API while trying to render "
         f"your graph after {max_retries} retries. "
     ) + error_msg_suffix
     raise ValueError(msg)

langchain_core/runnables/graph_png.py

@@ -1,36 +1,31 @@
 """Helper class to draw a state graph into a PNG file."""
 
-from typing import Any, Optional
+from typing import Any
 
 from langchain_core.runnables.graph import Graph, LabelsDict
 
+try:
+    import pygraphviz as pgv  # type: ignore[import-not-found]
+
+    _HAS_PYGRAPHVIZ = True
+except ImportError:
+    _HAS_PYGRAPHVIZ = False
+
 
 class PngDrawer:
     """Helper class to draw a state graph into a PNG file.
 
     It requires `graphviz` and `pygraphviz` to be installed.
-    :param fontname: The font to use for the labels
-    :param labels: A dictionary of label overrides. The dictionary
-        should have the following format:
-        {
-            "nodes": {
-                "node1": "CustomLabel1",
-                "node2": "CustomLabel2",
-                "__end__": "End Node"
-            },
-            "edges": {
-                "continue": "ContinueLabel",
-                "end": "EndLabel"
-            }
-        }
-        The keys are the original labels, and the values are the new labels.
-    Usage:
+
+    Example:
+        ```python
         drawer = PngDrawer()
-        drawer.draw(state_graph, 'graph.png')
+        drawer.draw(state_graph, "graph.png")
+        ```
     """
 
     def __init__(
-        self, fontname: Optional[str] = None, labels: Optional[LabelsDict] = None
+        self, fontname: str | None = None, labels: LabelsDict | None = None
     ) -> None:
         """Initializes the PNG drawer.
 
@@ -50,7 +45,7 @@ class PngDrawer:
                     }
                 }
                 The keys are the original labels, and the values are the new labels.
-                Defaults to None.
+
         """
         self.fontname = fontname or "arial"
         self.labels = labels or LabelsDict(nodes={}, edges={})
@@ -85,9 +80,6 @@ class PngDrawer:
         Args:
             viz: The graphviz object.
             node: The node to add.
-
-        Returns:
-            None
         """
         viz.add_node(
             node,
@@ -103,7 +95,7 @@ class PngDrawer:
         viz: Any,
         source: str,
         target: str,
-        label: Optional[str] = None,
+        label: str | None = None,
         conditional: bool = False,  # noqa: FBT001,FBT002
     ) -> None:
         """Adds an edge to the graph.
@@ -112,11 +104,8 @@ class PngDrawer:
             viz: The graphviz object.
             source: The source node.
             target: The target node.
-            label: The label for the edge. Defaults to None.
-            conditional: Whether the edge is conditional. Defaults to False.
-
-        Returns:
-            None
+            label: The label for the edge.
+            conditional: Whether the edge is conditional.
         """
         viz.add_edge(
             source,
@@ -127,18 +116,24 @@ class PngDrawer:
             style="dotted" if conditional else "solid",
         )
 
-    def draw(self, graph: Graph, output_path: Optional[str] = None) -> Optional[bytes]:
+    def draw(self, graph: Graph, output_path: str | None = None) -> bytes | None:
         """Draw the given state graph into a PNG file.
 
         Requires `graphviz` and `pygraphviz` to be installed.
-        :param graph: The graph to draw
-        :param output_path: The path to save the PNG. If None, PNG bytes are returned.
+
+        Args:
+            graph: The graph to draw
+            output_path: The path to save the PNG. If `None`, PNG bytes are returned.
+
+        Raises:
+            ImportError: If `pygraphviz` is not installed.
+
+        Returns:
+            The PNG bytes if `output_path` is None, else None.
         """
-        try:
-            import pygraphviz as pgv  # type: ignore[import-not-found]
-        except ImportError as exc:
+        if not _HAS_PYGRAPHVIZ:
             msg = "Install pygraphviz to draw graphs: `pip install pygraphviz`."
-            raise ImportError(msg) from exc
+            raise ImportError(msg)
 
         # Create a directed graph
         viz = pgv.AGraph(directed=True, nodesep=0.9, ranksep=1.0)