From e0da514bbc9808803d777aa15a4ea496bf3f1f84 Mon Sep 17 00:00:00 2001 From: FenjuFu Date: Mon, 27 Jul 2026 19:49:54 +0800 Subject: [PATCH] add Mermaid graph visualization --- burr/core/application.py | 10 +-- burr/core/graph.py | 85 +++++++++++++++++++++++-- docs/getting_started/simple-example.rst | 11 ++++ tests/core/test_graphviz_display.py | 29 +++++++++ tests/core/test_mermaid_display.py | 64 +++++++++++++++++++ 5 files changed, 188 insertions(+), 11 deletions(-) create mode 100644 tests/core/test_mermaid_display.py diff --git a/burr/core/application.py b/burr/core/application.py index 25bce4a10..decdb1d1b 100644 --- a/burr/core/application.py +++ b/burr/core/application.py @@ -2002,21 +2002,21 @@ def visualize( include_conditions: bool = False, include_state: bool = False, view: bool = False, - engine: Literal["graphviz"] = "graphviz", + engine: Literal["graphviz", "mermaid"] = "graphviz", write_dot: bool = False, **engine_kwargs: Any, - ) -> Optional["graphviz.Digraph"]: # noqa: F821 - """Visualizes the application graph using graphviz. This will render the graph. + ) -> Union["graphviz.Digraph", str, None]: # noqa: F821 + """Visualizes the application graph using Graphviz or Mermaid. :param output_file_path: The path to save this to, None if you don't want to save. Do not pass an extension for graphviz, instead pass `format` in `engine_kwargs` (e.g. `format="png"`) :param include_conditions: Whether to include condition strings on the edges (this can get noisy) :param include_state: Whether to indicate the action "signature" (reads/writes) on the nodes :param view: Whether to bring up a view - :param engine: The engine to use -- only graphviz is supported for now + :param engine: ``graphviz`` (default) or dependency-free Mermaid text :param write_dot: If True, produce a graphviz dot file :param engine_kwargs: Additional kwargs to pass to the engine - :return: The graphviz object + :return: A Graphviz object or Mermaid flowchart text """ return self.graph.visualize( output_file_path=output_file_path, diff --git a/burr/core/graph.py b/burr/core/graph.py index fcbc9c5cb..2ea856cb8 100644 --- a/burr/core/graph.py +++ b/burr/core/graph.py @@ -110,6 +110,63 @@ def _render_graphviz( pathlib.Path(f"{path_without_suffix}.{fmt}").write_bytes(graphviz_obj.pipe(format=fmt)) +def _escape_mermaid_label(value: str) -> str: + """Escape text embedded in a Mermaid quoted node or edge label.""" + return ( + value.replace("&", "&") + .replace('"', """) + .replace("<", "<") + .replace(">", ">") + .replace("\n", "
") + ) + + +def _render_mermaid( + actions: List[Action], + transitions: List[Transition], + include_conditions: bool, + include_state: bool, + direction: str = "TD", +) -> str: + """Render a Burr graph as dependency-free Mermaid flowchart text.""" + valid_directions = {"TB", "TD", "BT", "RL", "LR"} + if direction not in valid_directions: + raise ValueError( + f"Invalid Mermaid direction {direction!r}. Expected one of {sorted(valid_directions)}" + ) + + action_ids = {action.name: f"action_{index}" for index, action in enumerate(actions)} + input_ids = {} + lines = [f"flowchart {direction}"] + for action in actions: + label = ( + action.name + if not include_state + else f"{action.name}({', '.join(action.reads)}): {', '.join(action.writes)}" + ) + lines.append(f' {action_ids[action.name]}["{_escape_mermaid_label(label)}"]') + required_inputs, optional_inputs = action.optional_and_required_inputs + for input_ in sorted(required_inputs | optional_inputs): + if input_.startswith("__"): + continue + if input_ not in input_ids: + input_ids[input_] = f"input_{len(input_ids)}" + lines.append(f' {input_ids[input_]}["input: {_escape_mermaid_label(input_)}"]') + lines.append(f" {input_ids[input_]} -.-> {action_ids[action.name]}") + + for transition in transitions: + source = action_ids[transition.from_.name] + target = action_ids[transition.to.name] + if include_conditions and transition.condition is not default: + label = _escape_mermaid_label(transition.condition.name) + lines.append(f' {source} -. "{label}" .-> {target}') + elif transition.condition is not default: + lines.append(f" {source} -.-> {target}") + else: + lines.append(f" {source} --> {target}") + return "\n".join(lines) + "\n" + + @dataclasses.dataclass class Graph: """Graph class allows you to specify actions and transitions between them. @@ -183,24 +240,40 @@ def visualize( include_conditions: bool = False, include_state: bool = False, view: bool = False, - engine: Literal["graphviz"] = "graphviz", + engine: Literal["graphviz", "mermaid"] = "graphviz", write_dot: bool = False, **engine_kwargs: Any, - ) -> Optional["graphviz.Digraph"]: # noqa: F821 - """Visualizes the graph using graphviz. This will render the graph. + ) -> Union["graphviz.Digraph", str, None]: # noqa: F821 + """Visualizes the graph using Graphviz or Mermaid. :param output_file_path: The path to save this to, None if you don't want to save. Do not pass an extension for graphviz, instead pass `format` in `engine_kwargs` (e.g. `format="png"`) :param include_conditions: Whether to include condition strings on the edges (this can get noisy) :param include_state: Whether to indicate the action "signature" (reads/writes) on the nodes :param view: Whether to bring up a view - :param engine: The engine to use -- only graphviz is supported for now + :param engine: ``graphviz`` (default) or dependency-free Mermaid text :param write_dot: If True, produce a graphviz dot file :param engine_kwargs: Additional kwargs to pass to the engine - :return: The graphviz object + :return: A Graphviz object or Mermaid flowchart text """ + if engine == "mermaid": + if view or write_dot: + raise ValueError("Mermaid does not support view or write_dot") + direction = engine_kwargs.pop("direction", "TD") + if engine_kwargs: + raise ValueError(f"Unsupported Mermaid options: {', '.join(sorted(engine_kwargs))}") + diagram = _render_mermaid( + self.actions, + self.transitions, + include_conditions, + include_state, + direction, + ) + if output_file_path: + pathlib.Path(output_file_path).write_text(diagram, encoding="utf-8") + return diagram if engine != "graphviz": - raise ValueError(f"Only graphviz is supported for now, not {engine}") + raise ValueError(f"Unsupported visualization engine: {engine}") try: import graphviz # noqa: F401 except ModuleNotFoundError: diff --git a/docs/getting_started/simple-example.rst b/docs/getting_started/simple-example.rst index de467c68d..b2ba03a1d 100644 --- a/docs/getting_started/simple-example.rst +++ b/docs/getting_started/simple-example.rst @@ -128,6 +128,17 @@ We can visualize the application (note you need ``burr[graphviz]`` installed): app.visualize("./graph", format="png") +Mermaid text can be generated without installing Graphviz. It can be pasted into +Markdown renderers that support Mermaid or saved directly to a file: + +.. code-block:: python + + mermaid = app.visualize( + engine="mermaid", + include_conditions=True, + output_file_path="./graph.mmd", + ) + .. image:: ../_static/chatbot.png :align: center diff --git a/tests/core/test_graphviz_display.py b/tests/core/test_graphviz_display.py index 40c016089..dc8ba28f0 100644 --- a/tests/core/test_graphviz_display.py +++ b/tests/core/test_graphviz_display.py @@ -68,3 +68,32 @@ def test_visualize_no_dot_output(graph, tmp_path: pathlib.Path): graph.visualize(output_file_path=None) assert not dot_file_path.exists() + + +def test_visualize_mermaid(graph): + diagram = graph.visualize(engine="mermaid", include_conditions=True) + + assert diagram == ("flowchart TD\n" ' action_0["counter"]\n' " action_0 --> action_0\n") + + +def test_visualize_mermaid_file(graph, tmp_path: pathlib.Path): + output_file = tmp_path / "graph.mmd" + + diagram = graph.visualize( + engine="mermaid", + output_file_path=output_file, + direction="LR", + ) + + assert output_file.read_text(encoding="utf-8") == diagram + assert diagram.startswith("flowchart LR\n") + + +def test_visualize_mermaid_rejects_graphviz_options(graph): + with pytest.raises(ValueError, match="does not support"): + graph.visualize(engine="mermaid", view=True) + + +def test_visualize_mermaid_rejects_invalid_direction(graph): + with pytest.raises(ValueError, match="Invalid Mermaid direction"): + graph.visualize(engine="mermaid", direction="sideways") diff --git a/tests/core/test_mermaid_display.py b/tests/core/test_mermaid_display.py new file mode 100644 index 000000000..340b5fcb9 --- /dev/null +++ b/tests/core/test_mermaid_display.py @@ -0,0 +1,64 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +from burr.core import Condition +from burr.core.graph import GraphBuilder + +from tests.core.test_graph import PassedInAction + + +def _action(inputs=None): + return PassedInAction( + reads=["count"], + writes=["count"], + fn=lambda state: {"count": state.get("count", 0) + 1}, + update_fn=lambda result, state: state.update(**result), + inputs=inputs or [], + ) + + +def test_mermaid_includes_inputs_state_and_conditions(): + graph = ( + GraphBuilder() + .with_actions(counter=_action(["prompt"]), result=_action()) + .with_transitions( + ("counter", "result", Condition.expr("count < 10")), + ("result", "counter"), + ) + .build() + ) + + diagram = graph.visualize( + engine="mermaid", + include_conditions=True, + include_state=True, + ) + + assert 'action_0["counter(count): count"]' in diagram + assert 'input_0["input: prompt"]' in diagram + assert "input_0 -.-> action_0" in diagram + assert 'action_0 -. "count < 10" .-> action_1' in diagram + assert "action_1 --> action_0" in diagram + + +def test_mermaid_ignores_framework_injected_inputs(): + graph = GraphBuilder().with_actions(counter=_action(["__context", "prompt"])).build() + + diagram = graph.visualize(engine="mermaid") + + assert "__context" not in diagram + assert "input: prompt" in diagram