From e831751e3fc2ab5997a349d3f4e65ef59ab16930 Mon Sep 17 00:00:00 2001
From: goulustis <cnimaiv@gmail.com>
Date: Fri, 28 Nov 2025 18:17:40 +0800
Subject: [PATCH] return the image instead

---
 lang_agent/base.py | 10 +++++++---
 1 file changed, 7 insertions(+), 3 deletions(-)

diff --git a/lang_agent/base.py b/lang_agent/base.py
index c19be63..ba7c2d6 100644
--- a/lang_agent/base.py
+++ b/lang_agent/base.py
@@ -25,7 +25,7 @@ class GraphBase(ABC):
     def invoke(self, *nargs, **kwargs):
         pass
 
-    def show_graph(self):
+    def show_graph(self, ret_img:bool=False):
         #NOTE: just a useful tool for debugging; has zero useful functionality
         
         err_str = f"{type(self)} does not have workflow, this is unsupported"
@@ -33,8 +33,12 @@ class GraphBase(ABC):
 
         logger.info("creating image")
         img = Image.open(BytesIO(self.workflow.get_graph().draw_mermaid_png()))
-        plt.imshow(img)
-        plt.show()
+
+        if not ret_img:
+            plt.imshow(img)
+            plt.show()
+        else:
+            return img
 
 
 class ToolNodeBase(GraphBase):