langchain-core 1.0.0a6__py3-none-any.whl → 1.0.4__py3-none-any.whl

langchain_core/runnables/graph_mermaid.py

@@ -10,7 +10,7 @@ import string
 import time
 from dataclasses import asdict
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Literal, Optional
+from typing import TYPE_CHECKING, Any, Literal
 
 import yaml
 
@@ -45,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
@@ -164,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)
 
@@ -280,34 +274,28 @@ 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: Optional[str] = None,
+    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.
-        base_url (str, optional): Base URL for the Mermaid.ink API.
-            Defaults to None.
+        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.
@@ -340,8 +328,8 @@ 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:
@@ -412,12 +400,12 @@ 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: Optional[str] = None,
+    base_url: str | None = None,
 ) -> bytes:
     """Renders Mermaid graph using the Mermaid.INK API."""
     # Defaults to using the public mermaid.ink server.
@@ -466,7 +454,10 @@ def _render_mermaid_using_api(
                 return img_bytes
 
             # If we get a server error (5xx), retry
-            if 500 <= response.status_code < 600 and attempt < max_retries:
+            if (
+                requests.codes.internal_server_error <= response.status_code
+                and attempt < max_retries
+            ):
                 # Exponential backoff with jitter
                 sleep_time = retry_delay * (2**attempt) * (0.5 + 0.5 * random.random())  # noqa: S311 not used for crypto
                 time.sleep(sleep_time)

langchain_core/runnables/graph_png.py

@@ -1,6 +1,7 @@
 """Helper class to draw a state graph into a PNG file."""
 
-from typing import Any, Optional
+from itertools import groupby
+from typing import Any
 
 from langchain_core.runnables.graph import Graph, LabelsDict
 
@@ -15,18 +16,17 @@ except ImportError:
 class PngDrawer:
     """Helper class to draw a state graph into a PNG file.
 
-    It requires ``graphviz`` and ``pygraphviz`` to be installed.
+    It requires `graphviz` and `pygraphviz` to be installed.
 
     Example:
-
-        .. code-block:: python
-
-            drawer = PngDrawer()
-            drawer.draw(state_graph, "graph.png")
+        ```python
+        drawer = PngDrawer()
+        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.
 
@@ -46,7 +46,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={})
@@ -96,7 +96,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.
@@ -105,8 +105,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.
+            label: The label for the edge.
+            conditional: Whether the edge is conditional.
         """
         viz.add_edge(
             source,
@@ -117,20 +117,20 @@ 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.
 
         Args:
             graph: The graph to draw
-            output_path: The path to save the PNG. If None, PNG bytes are returned.
+            output_path: The path to save the PNG. If `None`, PNG bytes are returned.
 
         Raises:
-            ImportError: If ``pygraphviz`` is not installed.
+            ImportError: If `pygraphviz` is not installed.
 
         Returns:
-            The PNG bytes if ``output_path`` is None, else None.
+            The PNG bytes if `output_path` is None, else None.
         """
         if not _HAS_PYGRAPHVIZ:
             msg = "Install pygraphviz to draw graphs: `pip install pygraphviz`."
@@ -142,6 +142,7 @@ class PngDrawer:
         # Add nodes, conditional edges, and edges to the graph
         self.add_nodes(viz, graph)
         self.add_edges(viz, graph)
+        self.add_subgraph(viz, [node.split(":") for node in graph.nodes])
 
         # Update entrypoint and END styles
         self.update_styles(viz, graph)
@@ -162,6 +163,32 @@ class PngDrawer:
         for node in graph.nodes:
             self.add_node(viz, node)
 
+    def add_subgraph(
+        self,
+        viz: Any,
+        nodes: list[list[str]],
+        parent_prefix: list[str] | None = None,
+    ) -> None:
+        """Add subgraphs to the graph.
+
+        Args:
+            viz: The graphviz object.
+            nodes: The nodes to add.
+            parent_prefix: The prefix of the parent subgraph.
+        """
+        for prefix, grouped in groupby(
+            [node[:] for node in sorted(nodes)],
+            key=lambda x: x.pop(0),
+        ):
+            current_prefix = (parent_prefix or []) + [prefix]
+            grouped_nodes = list(grouped)
+            if len(grouped_nodes) > 1:
+                subgraph = viz.add_subgraph(
+                    [":".join(current_prefix + node) for node in grouped_nodes],
+                    name="cluster_" + ":".join(current_prefix),
+                )
+                self.add_subgraph(subgraph, grouped_nodes, current_prefix)
+
     def add_edges(self, viz: Any, graph: Graph) -> None:
         """Add edges to the graph.