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

langchain_core/runnables/graph.py

@@ -4,21 +4,21 @@ from __future__ import annotations
 
 import inspect
 from collections import defaultdict
+from collections.abc import Callable
 from dataclasses import dataclass, field
 from enum import Enum
 from typing import (
     TYPE_CHECKING,
     Any,
-    Callable,
     NamedTuple,
-    Optional,
     Protocol,
     TypedDict,
-    Union,
     overload,
 )
 from uuid import UUID, uuid4
 
+from langchain_core.load.serializable import to_json_not_implemented
+from langchain_core.runnables.base import Runnable, RunnableSerializable
 from langchain_core.utils.pydantic import _IgnoreUnserializable, is_basemodel_subclass
 
 if TYPE_CHECKING:
@@ -52,7 +52,7 @@ def is_uuid(value: str) -> bool:
         value: The string to check.
 
     Returns:
-        True if the string is a valid UUID, False otherwise.
+        `True` if the string is a valid UUID, `False` otherwise.
     """
     try:
         UUID(value)
@@ -62,28 +62,23 @@ def is_uuid(value: str) -> bool:
 
 
 class Edge(NamedTuple):
-    """Edge in a graph.
-
-    Parameters:
-        source: The source node id.
-        target: The target node id.
-        data: Optional data associated with the edge. Defaults to None.
-        conditional: Whether the edge is conditional. Defaults to False.
-    """
+    """Edge in a graph."""
 
     source: str
+    """The source node id."""
     target: str
-    data: Optional[Stringifiable] = None
+    """The target node id."""
+    data: Stringifiable | None = None
+    """Optional data associated with the edge. """
     conditional: bool = False
+    """Whether the edge is conditional."""
 
-    def copy(
-        self, *, source: Optional[str] = None, target: Optional[str] = None
-    ) -> Edge:
+    def copy(self, *, source: str | None = None, target: str | None = None) -> Edge:
         """Return a copy of the edge with optional new source and target nodes.
 
         Args:
-            source: The new source node id. Defaults to None.
-            target: The new target node id. Defaults to None.
+            source: The new source node id.
+            target: The new target node id.
 
         Returns:
             A copy of the edge with the new source and target nodes.
@@ -97,31 +92,28 @@ class Edge(NamedTuple):
 
 
 class Node(NamedTuple):
-    """Node in a graph.
-
-    Parameters:
-        id: The unique identifier of the node.
-        name: The name of the node.
-        data: The data of the node.
-        metadata: Optional metadata for the node. Defaults to None.
-    """
+    """Node in a graph."""
 
     id: str
+    """The unique identifier of the node."""
     name: str
-    data: Union[type[BaseModel], RunnableType, None]
-    metadata: Optional[dict[str, Any]]
+    """The name of the node."""
+    data: type[BaseModel] | RunnableType | None
+    """The data of the node."""
+    metadata: dict[str, Any] | None
+    """Optional metadata for the node. """
 
     def copy(
         self,
         *,
-        id: Optional[str] = None,
-        name: Optional[str] = None,
+        id: str | None = None,
+        name: str | None = None,
     ) -> Node:
         """Return a copy of the node with optional new id and name.
 
         Args:
-            id: The new node id. Defaults to None.
-            name: The new node name. Defaults to None.
+            id: The new node id.
+            name: The new node name.
 
         Returns:
             A copy of the node with the new id and name.
@@ -135,16 +127,12 @@ class Node(NamedTuple):
 
 
 class Branch(NamedTuple):
-    """Branch in a graph.
-
-    Parameters:
-        condition: A callable that returns a string representation of the condition.
-        ends: Optional dictionary of end node ids for the branches. Defaults
-            to None.
-    """
+    """Branch in a graph."""
 
     condition: Callable[..., str]
-    ends: Optional[dict[str, str]]
+    """A callable that returns a string representation of the condition."""
+    ends: dict[str, str] | None
+    """Optional dictionary of end node ids for the branches. """
 
 
 class CurveStyle(Enum):
@@ -168,10 +156,10 @@ class CurveStyle(Enum):
 class NodeStyles:
     """Schema for Hexadecimal color codes for different node types.
 
-    Parameters:
-        default: The default color code. Defaults to "fill:#f2f0ff,line-height:1.2".
-        first: The color code for the first node. Defaults to "fill-opacity:0".
-        last: The color code for the last node. Defaults to "fill:#bfb6fc".
+    Args:
+        default: The default color code.
+        first: The color code for the first node.
+        last: The color code for the last node.
     """
 
     default: str = "fill:#f2f0ff,line-height:1.2"
@@ -182,13 +170,15 @@ class NodeStyles:
 class MermaidDrawMethod(Enum):
     """Enum for different draw methods supported by Mermaid."""
 
-    PYPPETEER = "pyppeteer"  # Uses Pyppeteer to render the graph
-    API = "api"  # Uses Mermaid.INK API to render the graph
+    PYPPETEER = "pyppeteer"
+    """Uses Pyppeteer to render the graph"""
+    API = "api"
+    """Uses Mermaid.INK API to render the graph"""
 
 
 def node_data_str(
     id: str,
-    data: Union[type[BaseModel], RunnableType, None],
+    data: type[BaseModel] | RunnableType | None,
 ) -> str:
     """Convert the data of a node to a string.
 
@@ -199,8 +189,6 @@ def node_data_str(
     Returns:
         A string representation of the data.
     """
-    from langchain_core.runnables.base import Runnable
-
     if not is_uuid(id) or data is None:
         return id
     data_str = data.get_name() if isinstance(data, Runnable) else data.__name__
@@ -209,20 +197,17 @@ def node_data_str(
 
 def node_data_json(
     node: Node, *, with_schemas: bool = False
-) -> dict[str, Union[str, dict[str, Any]]]:
+) -> dict[str, str | dict[str, Any]]:
     """Convert the data of a node to a JSON-serializable format.
 
     Args:
-        node: The node to convert.
-        with_schemas: Whether to include the schema of the data if
-            it is a Pydantic model. Defaults to False.
+        node: The `Node` to convert.
+        with_schemas: Whether to include the schema of the data if it is a Pydantic
+            model.
 
     Returns:
         A dictionary with the type of the data and the data itself.
     """
-    from langchain_core.load.serializable import to_json_not_implemented
-    from langchain_core.runnables.base import Runnable, RunnableSerializable
-
     if node.data is None:
         json: dict[str, Any] = {}
     elif isinstance(node.data, RunnableSerializable):
@@ -269,7 +254,7 @@ def node_data_json(
 class Graph:
     """Graph of nodes and edges.
 
-    Parameters:
+    Args:
         nodes: Dictionary of nodes in the graph. Defaults to an empty dictionary.
         edges: List of edges in the graph. Defaults to an empty list.
     """
@@ -282,7 +267,7 @@ class Graph:
 
         Args:
             with_schemas: Whether to include the schemas of the nodes if they are
-                Pydantic models. Defaults to False.
+                Pydantic models.
 
         Returns:
             A dictionary with the nodes and edges of the graph.
@@ -327,17 +312,17 @@ class Graph:
 
     def add_node(
         self,
-        data: Union[type[BaseModel], RunnableType, None],
-        id: Optional[str] = None,
+        data: type[BaseModel] | RunnableType | None,
+        id: str | None = None,
         *,
-        metadata: Optional[dict[str, Any]] = None,
+        metadata: dict[str, Any] | None = None,
     ) -> Node:
         """Add a node to the graph and return it.
 
         Args:
             data: The data of the node.
-            id: The id of the node. Defaults to None.
-            metadata: Optional metadata for the node. Defaults to None.
+            id: The id of the node.
+            metadata: Optional metadata for the node.
 
         Returns:
             The node that was added to the graph.
@@ -368,7 +353,7 @@ class Graph:
         self,
         source: Node,
         target: Node,
-        data: Optional[Stringifiable] = None,
+        data: Stringifiable | None = None,
         conditional: bool = False,  # noqa: FBT001,FBT002
     ) -> Edge:
         """Add an edge to the graph and return it.
@@ -376,8 +361,8 @@ class Graph:
         Args:
             source: The source node of the edge.
             target: The target node of the edge.
-            data: Optional data associated with the edge. Defaults to None.
-            conditional: Whether the edge is conditional. Defaults to False.
+            data: Optional data associated with the edge.
+            conditional: Whether the edge is conditional.
 
         Returns:
             The edge that was added to the graph.
@@ -399,14 +384,14 @@ class Graph:
 
     def extend(
         self, graph: Graph, *, prefix: str = ""
-    ) -> tuple[Optional[Node], Optional[Node]]:
+    ) -> tuple[Node | None, Node | None]:
         """Add all nodes and edges from another graph.
 
         Note this doesn't check for duplicates, nor does it connect the graphs.
 
         Args:
             graph: The graph to add.
-            prefix: The prefix to add to the node ids. Defaults to "".
+            prefix: The prefix to add to the node ids.
 
         Returns:
             A tuple of the first and last nodes of the subgraph.
@@ -470,19 +455,27 @@ class Graph:
             ],
         )
 
-    def first_node(self) -> Optional[Node]:
+    def first_node(self) -> Node | None:
         """Find the single node that is not a target of any edge.
 
-        If there is no such node, or there are multiple, return None.
+        If there is no such node, or there are multiple, return `None`.
         When drawing the graph, this node would be the origin.
+
+        Returns:
+            The first node, or None if there is no such node or multiple
+            candidates.
         """
         return _first_node(self)
 
-    def last_node(self) -> Optional[Node]:
+    def last_node(self) -> Node | None:
         """Find the single node that is not a source of any edge.
 
-        If there is no such node, or there are multiple, return None.
+        If there is no such node, or there are multiple, return `None`.
         When drawing the graph, this node would be the destination.
+
+        Returns:
+            The last node, or None if there is no such node or multiple
+            candidates.
         """
         return _last_node(self)
 
@@ -513,8 +506,13 @@ class Graph:
             self.remove_node(last_node)
 
     def draw_ascii(self) -> str:
-        """Draw the graph as an ASCII art string."""
-        from langchain_core.runnables.graph_ascii import draw_ascii
+        """Draw the graph as an ASCII art string.
+
+        Returns:
+            The ASCII art string.
+        """
+        # Import locally to prevent circular import
+        from langchain_core.runnables.graph_ascii import draw_ascii  # noqa: PLC0415
 
         return draw_ascii(
             {node.id: node.name for node in self.nodes.values()},
@@ -529,36 +527,38 @@ class Graph:
     def draw_png(
         self,
         output_file_path: str,
-        fontname: Optional[str] = None,
-        labels: Optional[LabelsDict] = None,
+        fontname: str | None = None,
+        labels: LabelsDict | None = None,
     ) -> None: ...
 
     @overload
     def draw_png(
         self,
         output_file_path: None,
-        fontname: Optional[str] = None,
-        labels: Optional[LabelsDict] = None,
+        fontname: str | None = None,
+        labels: LabelsDict | None = None,
     ) -> bytes: ...
 
     def draw_png(
         self,
-        output_file_path: Optional[str] = None,
-        fontname: Optional[str] = None,
-        labels: Optional[LabelsDict] = None,
-    ) -> Union[bytes, None]:
+        output_file_path: str | None = None,
+        fontname: str | None = None,
+        labels: LabelsDict | None = None,
+    ) -> bytes | None:
         """Draw the graph as a PNG image.
 
         Args:
-            output_file_path: The path to save the image to. If None, the image
-                is not saved. Defaults to None.
-            fontname: The name of the font to use. Defaults to None.
-            labels: Optional labels for nodes and edges in the graph. Defaults to None.
+            output_file_path: The path to save the image to. If `None`, the image
+                is not saved.
+            fontname: The name of the font to use.
+            labels: Optional labels for nodes and edges in the graph. Defaults to
+                `None`.
 
         Returns:
             The PNG image as bytes if output_file_path is None, None otherwise.
         """
-        from langchain_core.runnables.graph_png import PngDrawer
+        # Import locally to prevent circular import
+        from langchain_core.runnables.graph_png import PngDrawer  # noqa: PLC0415
 
         default_node_labels = {node.id: node.name for node in self.nodes.values()}
 
@@ -578,42 +578,39 @@ class Graph:
         *,
         with_styles: bool = True,
         curve_style: CurveStyle = CurveStyle.LINEAR,
-        node_colors: Optional[NodeStyles] = None,
+        node_colors: NodeStyles | None = None,
         wrap_label_n_words: int = 9,
-        frontmatter_config: Optional[dict[str, Any]] = None,
+        frontmatter_config: dict[str, Any] | None = None,
     ) -> str:
         """Draw the graph as a Mermaid syntax string.
 
         Args:
-            with_styles: Whether to include styles in the syntax. Defaults to True.
-            curve_style: The style of the edges. Defaults to CurveStyle.LINEAR.
-            node_colors: The colors of the nodes. Defaults to NodeStyles().
+            with_styles: Whether to include styles in the syntax.
+            curve_style: The style of the edges.
+            node_colors: The colors of the nodes.
             wrap_label_n_words: The number of words to wrap the node labels at.
-                Defaults to 9.
-            frontmatter_config (dict[str, Any], optional): Mermaid frontmatter config.
+            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:
             The Mermaid syntax string.
-
         """
-        from langchain_core.runnables.graph_mermaid import draw_mermaid
+        # Import locally to prevent circular import
+        from langchain_core.runnables.graph_mermaid import draw_mermaid  # noqa: PLC0415
 
         graph = self.reid()
         first_node = graph.first_node()
@@ -635,56 +632,58 @@ class Graph:
         self,
         *,
         curve_style: CurveStyle = CurveStyle.LINEAR,
-        node_colors: Optional[NodeStyles] = None,
+        node_colors: NodeStyles | None = None,
         wrap_label_n_words: int = 9,
-        output_file_path: Optional[str] = None,
+        output_file_path: str | None = None,
         draw_method: MermaidDrawMethod = MermaidDrawMethod.API,
         background_color: str = "white",
         padding: int = 10,
         max_retries: int = 1,
         retry_delay: float = 1.0,
-        frontmatter_config: Optional[dict[str, Any]] = None,
+        frontmatter_config: dict[str, Any] | None = None,
+        base_url: str | None = None,
     ) -> bytes:
         """Draw the graph as a PNG image using Mermaid.
 
         Args:
-            curve_style: The style of the edges. Defaults to CurveStyle.LINEAR.
-            node_colors: The colors of the nodes. Defaults to NodeStyles().
+            curve_style: The style of the edges.
+            node_colors: The colors of the nodes.
             wrap_label_n_words: The number of words to wrap the node labels at.
-                Defaults to 9.
-            output_file_path: The path to save the image to. If None, the image
-                is not saved. Defaults to None.
+            output_file_path: The path to save the image to. If `None`, the image
+                is not saved.
             draw_method: The method to use to draw the graph.
-                Defaults to MermaidDrawMethod.API.
-            background_color: The color of the background. Defaults to "white".
-            padding: The padding around the graph. Defaults to 10.
-            max_retries: The maximum number of retries (MermaidDrawMethod.API).
-                Defaults to 1.
-            retry_delay: The delay between retries (MermaidDrawMethod.API).
-                Defaults to 1.0.
-            frontmatter_config (dict[str, Any], optional): Mermaid frontmatter config.
+            background_color: The color of the background.
+            padding: The padding around the graph.
+            max_retries: The maximum number of retries (`MermaidDrawMethod.API`).
+            retry_delay: The delay between retries (`MermaidDrawMethod.API`).
+            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"},
                     }
                 }
+                ```
+            base_url: The base URL of the Mermaid server for rendering via API.
+
 
         Returns:
             The PNG image as bytes.
 
         """
-        from langchain_core.runnables.graph_mermaid import draw_mermaid_png
+        # Import locally to prevent circular import
+        from langchain_core.runnables.graph_mermaid import (  # noqa: PLC0415
+            draw_mermaid_png,
+        )
 
         mermaid_syntax = self.draw_mermaid(
             curve_style=curve_style,
@@ -700,14 +699,15 @@ class Graph:
             padding=padding,
             max_retries=max_retries,
             retry_delay=retry_delay,
+            base_url=base_url,
         )
 
 
-def _first_node(graph: Graph, exclude: Sequence[str] = ()) -> Optional[Node]:
+def _first_node(graph: Graph, exclude: Sequence[str] = ()) -> Node | None:
     """Find the single node that is not a target of any edge.
 
     Exclude nodes/sources with ids in the exclude list.
-    If there is no such node, or there are multiple, return None.
+    If there is no such node, or there are multiple, return `None`.
     When drawing the graph, this node would be the origin.
     """
     targets = {edge.target for edge in graph.edges if edge.source not in exclude}
@@ -719,11 +719,11 @@ def _first_node(graph: Graph, exclude: Sequence[str] = ()) -> Optional[Node]:
     return found[0] if len(found) == 1 else None
 
 
-def _last_node(graph: Graph, exclude: Sequence[str] = ()) -> Optional[Node]:
+def _last_node(graph: Graph, exclude: Sequence[str] = ()) -> Node | None:
     """Find the single node that is not a source of any edge.
 
     Exclude nodes/targets with ids in the exclude list.
-    If there is no such node, or there are multiple, return None.
+    If there is no such node, or there are multiple, return `None`.
     When drawing the graph, this node would be the destination.
     """
     sources = {edge.source for edge in graph.edges if edge.target not in exclude}