mirror of https://github.com/hwchase17/langchain
Add PNG drawer for Runnable.get_graph() (#18239)
Thank you for contributing to LangChain!
- [ ] **PR title**: "package: description"
- Where "package" is whichever of langchain, community, core,
experimental, etc. is being modified. Use "docs: ..." for purely docs
changes, "templates: ..." for template changes, "infra: ..." for CI
changes.
- Example: "community: add foobar LLM"
- [ ] **PR message**: ***Delete this entire checklist*** and replace
with
- **Description:** a description of the change
- **Issue:** the issue # it fixes, if applicable
- **Dependencies:** any dependencies required for this change
- **Twitter handle:** if your PR gets announced, and you'd like a
mention, we'll gladly shout you out!
- [ ] **Add tests and docs**: If you're adding a new integration, please
include
1. a test for the integration, preferably unit tests that do not rely on
network access,
2. an example notebook showing its use. It lives in
`docs/docs/integrations` directory.
- [ ] **Lint and test**: Run `make format`, `make lint` and `make test`
from the root of the package(s) you've modified. See contribution
guidelines for more: https://python.langchain.com/docs/contributing/
Additional guidelines:
- Make sure optional dependencies are imported within a function.
- Please do not add dependencies to pyproject.toml files (even optional
ones) unless they are required for unit tests.
- Most PRs should not touch more than one package.
- Changes should be backwards compatible.
- If you are adding something to community, do not re-import it in
langchain.
If no one reviews your PR within a few days, please @-mention one of
baskaryan, efriis, eyurtsev, hwchase17.
pull/18077/head^2
Nuno Campos
7 months ago
committed by
GitHub
parent
d9fd1194f5
commit
6da08d0f22
@ -0,0 +1,108 @@
|
||||
from typing import Any, Optional
|
||||
|
||||
from langchain_core.runnables.graph import Graph, LabelsDict
|
||||
|
||||
|
||||
class PngDrawer:
|
||||
"""
|
||||
A helper class to draw a state graph into a PNG file.
|
||||
Requires graphviz and pygraphviz to be installed.
|
||||
:param fontname: The font to use for the labels
|
||||
:param 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.
|
||||
Usage:
|
||||
drawer = PngDrawer()
|
||||
drawer.draw(state_graph, 'graph.png')
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self, fontname: Optional[str] = None, labels: Optional[LabelsDict] = None
|
||||
) -> None:
|
||||
self.fontname = fontname or "arial"
|
||||
self.labels = labels or LabelsDict(nodes={}, edges={})
|
||||
|
||||
def get_node_label(self, label: str) -> str:
|
||||
label = self.labels.get("nodes", {}).get(label, label)
|
||||
return f"<<B>{label}</B>>"
|
||||
|
||||
def get_edge_label(self, label: str) -> str:
|
||||
label = self.labels.get("edges", {}).get(label, label)
|
||||
return f"<<U>{label}</U>>"
|
||||
|
||||
def add_node(self, viz: Any, node: str) -> None:
|
||||
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, source: str, target: str, label: Optional[str] = None
|
||||
) -> None:
|
||||
viz.add_edge(
|
||||
source,
|
||||
target,
|
||||
label=self.get_edge_label(label) if label else "",
|
||||
fontsize=12,
|
||||
fontname=self.fontname,
|
||||
)
|
||||
|
||||
def draw(self, graph: Graph, output_path: Optional[str] = None) -> Optional[bytes]:
|
||||
"""
|
||||
Draws the given state graph into a PNG file.
|
||||
Requires graphviz and pygraphviz to be installed.
|
||||
:param graph: The graph to draw
|
||||
:param output_path: The path to save the PNG. If None, PNG bytes are returned.
|
||||
"""
|
||||
|
||||
try:
|
||||
import pygraphviz as pgv # type: ignore[import]
|
||||
except ImportError as exc:
|
||||
raise ImportError(
|
||||
"Install pygraphviz to draw graphs: `pip install pygraphviz`."
|
||||
) from exc
|
||||
|
||||
# Create a directed graph
|
||||
viz = pgv.AGraph(directed=True, nodesep=0.9, ranksep=1.0)
|
||||
|
||||
# Add nodes, conditional edges, and edges to the graph
|
||||
self.add_nodes(viz, graph)
|
||||
self.add_edges(viz, graph)
|
||||
|
||||
# Update entrypoint and END styles
|
||||
self.update_styles(viz, graph)
|
||||
|
||||
# Save the graph as PNG
|
||||
try:
|
||||
return viz.draw(output_path, format="png", prog="dot")
|
||||
finally:
|
||||
viz.close()
|
||||
|
||||
def add_nodes(self, viz: Any, graph: Graph) -> None:
|
||||
for node in graph.nodes:
|
||||
self.add_node(viz, node)
|
||||
|
||||
def add_edges(self, viz: Any, graph: Graph) -> None:
|
||||
for start, end, label in graph.edges:
|
||||
self.add_edge(viz, start, end, label)
|
||||
|
||||
def update_styles(self, viz: Any, graph: Graph) -> None:
|
||||
if first := graph.first_node():
|
||||
viz.get_node(first.id).attr.update(fillcolor="lightblue")
|
||||
if last := graph.last_node():
|
||||
viz.get_node(last.id).attr.update(fillcolor="orange")
|
||||
Loading…
Reference in New Issue