feat: Make CompiledGraph displayable in Juypyter with display() (#2159)

This commit is contained in:
Christopher Brooks
2024-11-19 11:18:53 -05:00
parent e3a30a9b69
commit ca7da2fc41
+8
View File
@@ -628,3 +628,11 @@ class CompiledGraph(Pregel):
add_edge(key, end, conditional=True)
return graph
def _repr_mimebundle_(self, **kwargs) -> dict[str, Any]:
"""Mime bundle used by Jupyter to display the graph"""
output = {
"text/plain": repr(self),
"image/png": self.get_graph().draw_mermaid_png()
}
return output