PngDrawer Class — langchain Architecture

Source Code

libs/core/langchain_core/runnables/graph_png.py lines 16–215

class PngDrawer:
    """Helper class to draw a state graph into a PNG file.

    It requires `graphviz` and `pygraphviz` to be installed.

    Example:
        ```python
        drawer = PngDrawer()
        drawer.draw(state_graph, "graph.png")
        ```
    """

    def __init__(
        self, fontname: str | None = None, labels: LabelsDict | None = None
    ) -> None:
        """Initializes the PNG drawer.

        Args:
            fontname: The font to use for the labels. Defaults to "arial".
            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.

        """
        self.fontname = fontname or "arial"
        self.labels = labels or LabelsDict(nodes={}, edges={})

    def get_node_label(self, label: str) -> str:
        """Returns the label to use for a node.

        Args:
            label: The original label.

        Returns:
            The new label.
        """
        label = self.labels.get("nodes", {}).get(label, label)
        return f"<<B>{label}</B>>"

    def get_edge_label(self, label: str) -> str:
        """Returns the label to use for an edge.

        Args:
            label: The original label.

        Returns:
            The new label.
        """
        label = self.labels.get("edges", {}).get(label, label)
        return f"<<U>{label}</U>>"

    def add_node(self, viz: Any, node: str) -> None:
        """Adds a node to the graph.

        Args:
            viz: The graphviz object.
            node: The node to add.
        """
        viz.add_node(
            node,
            label=self.get_node_label(node),
            style="filled",
            fillcolor="yellow",
            fontsize=15,
            fontname=self.fontname,
        )

    def add_edge(
        self,
        viz: Any,