moved visualization to base

This commit is contained in:
2025-11-21 14:22:01 +08:00
parent 643f3a9dc3
commit 9022540b79
2 changed files with 18 additions and 15 deletions

View File

@@ -3,9 +3,6 @@ from typing import Type, TypedDict, Literal, Dict, List, Tuple, Any
import tyro
from pydantic import BaseModel, Field
from loguru import logger
from PIL import Image
from io import BytesIO
import matplotlib.pyplot as plt
import jax
import os.path as osp
import commentjson
@@ -254,12 +251,6 @@ class RoutingGraph(GraphBase):
return workflow
def show_graph(self):
logger.info("creating image")
img = Image.open(BytesIO(self.workflow.get_graph().draw_mermaid_png()))
plt.imshow(img)
plt.show()
if __name__ == "__main__":
from dotenv import load_dotenv
from langchain.messages import SystemMessage, HumanMessage