From 31a641a1553a5d06fdf976489fb11da3a9b34e75 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=C3=81ngel=20Igareta?= <32129522+angeligareta@users.noreply.github.com> Date: Wed, 3 Apr 2024 15:20:35 +0200 Subject: [PATCH] core: fix return of draw_mermaid_png and change to not save image by default (#19950) - **Description:** Improvement for #19599: fixing missing return of graph.draw_mermaid_png and improve it to make the saving of the rendered image optional Co-authored-by: Angel Igareta --- libs/core/langchain_core/runnables/graph.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/libs/core/langchain_core/runnables/graph.py b/libs/core/langchain_core/runnables/graph.py index 84ae717d03..8c11cd8fc7 100644 --- a/libs/core/langchain_core/runnables/graph.py +++ b/libs/core/langchain_core/runnables/graph.py @@ -369,11 +369,11 @@ class Graph: start="#ffdfba", end="#baffc9", other="#fad7de" ), wrap_label_n_words: int = 9, - output_file_path: str = "graph.png", + output_file_path: Optional[str] = None, draw_method: MermaidDrawMethod = MermaidDrawMethod.API, background_color: str = "white", padding: int = 10, - ) -> None: + ) -> bytes: from langchain_core.runnables.graph_mermaid import draw_mermaid_png mermaid_syntax = self.draw_mermaid( @@ -381,7 +381,7 @@ class Graph: node_colors=node_colors, wrap_label_n_words=wrap_label_n_words, ) - draw_mermaid_png( + return draw_mermaid_png( mermaid_syntax=mermaid_syntax, output_file_path=output_file_path, draw_method=draw_method,