Skip to content

structure_visualizer

StructureVisualizer

Utility class to visualize a Structure structure.

Source code in griptape/utils/structure_visualizer.py 
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
@define
class StructureVisualizer:
    """Utility class to visualize a Structure structure."""

    structure: Structure = field()
    header: str = field(default="graph TD;", kw_only=True)
    build_node_id: Callable[[BaseTask], str] = field(default=lambda task: task.id.title(), kw_only=True)
    query_params: dict[str, str] = field(factory=dict, kw_only=True)
    base_url: str = field(default="https://mermaid.ink", kw_only=True)

    def to_url(self) -> str:
        """Generates a url that renders the Workflow structure as a Mermaid flowchart.

        Reference: https://mermaid.js.org/ecosystem/tutorials#jupyter-integration-with-mermaid-js.

        Returns:
            str: URL to the rendered image
        """
        self.structure.resolve_relationships()

        tasks = self.__render_tasks(self.structure.tasks)
        graph = f"{self.header}{tasks}"

        graph_bytes = graph.encode("utf-8")
        base64_string = base64.b64encode(graph_bytes).decode("utf-8")

        url = urllib.parse.urljoin(self.base_url, f"svg/{base64_string}")
        req = PreparedRequest()
        req.prepare_url(url, self.query_params)

        if req.url is None:
            raise ValueError("Failed to generate the URL")

        return req.url

    def __render_tasks(self, tasks: list[BaseTask]) -> str:
        return "\n\t" + "\n\t".join([self.__render_task(task) for task in tasks])

    def __render_task(self, task: BaseTask) -> str:
        from griptape.drivers.structure_run.local import LocalStructureRunDriver
        from griptape.tasks import BranchTask, StructureRunTask

        parts = []
        if task.children:
            children = " & ".join([f"{self.build_node_id(child)}" for child in task.children])
            if isinstance(task, BranchTask):
                parts.append(f"{self.build_node_id(task)}{{ {self.build_node_id(task)} }}-.-> {children};")
            else:
                parts.append(f"{self.build_node_id(task)}--> {children};")
        else:
            parts.append(f"{self.build_node_id(task)};")

        if isinstance(task, StructureRunTask) and isinstance(task.structure_run_driver, LocalStructureRunDriver):
            sub_structure = task.structure_run_driver.create_structure()
            sub_tasks = self.__render_tasks(sub_structure.tasks)
            parts.append(f"subgraph {self.build_node_id(task)}{sub_tasks}\n\tend")

        return "\n\t".join(parts)

base_url = field(default='https://mermaid.ink', kw_only=True) class-attribute instance-attribute

build_node_id = field(default=lambda task: task.id.title(), kw_only=True) class-attribute instance-attribute

header = field(default='graph TD;', kw_only=True) class-attribute instance-attribute

query_params = field(factory=dict, kw_only=True) class-attribute instance-attribute

structure = field() class-attribute instance-attribute

__render_task(task)

Source code in griptape/utils/structure_visualizer.py 
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
def __render_task(self, task: BaseTask) -> str:
    from griptape.drivers.structure_run.local import LocalStructureRunDriver
    from griptape.tasks import BranchTask, StructureRunTask

    parts = []
    if task.children:
        children = " & ".join([f"{self.build_node_id(child)}" for child in task.children])
        if isinstance(task, BranchTask):
            parts.append(f"{self.build_node_id(task)}{{ {self.build_node_id(task)} }}-.-> {children};")
        else:
            parts.append(f"{self.build_node_id(task)}--> {children};")
    else:
        parts.append(f"{self.build_node_id(task)};")

    if isinstance(task, StructureRunTask) and isinstance(task.structure_run_driver, LocalStructureRunDriver):
        sub_structure = task.structure_run_driver.create_structure()
        sub_tasks = self.__render_tasks(sub_structure.tasks)
        parts.append(f"subgraph {self.build_node_id(task)}{sub_tasks}\n\tend")

    return "\n\t".join(parts)

__render_tasks(tasks)

Source code in griptape/utils/structure_visualizer.py 
50
51
def __render_tasks(self, tasks: list[BaseTask]) -> str:
    return "\n\t" + "\n\t".join([self.__render_task(task) for task in tasks])

to_url()

Generates a url that renders the Workflow structure as a Mermaid flowchart.

Reference: https://mermaid.js.org/ecosystem/tutorials#jupyter-integration-with-mermaid-js.

Returns:

Name Type Description
str str

URL to the rendered image
Source code in griptape/utils/structure_visualizer.py 
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
def to_url(self) -> str:
    """Generates a url that renders the Workflow structure as a Mermaid flowchart.

    Reference: https://mermaid.js.org/ecosystem/tutorials#jupyter-integration-with-mermaid-js.

    Returns:
        str: URL to the rendered image
    """
    self.structure.resolve_relationships()

    tasks = self.__render_tasks(self.structure.tasks)
    graph = f"{self.header}{tasks}"

    graph_bytes = graph.encode("utf-8")
    base64_string = base64.b64encode(graph_bytes).decode("utf-8")

    url = urllib.parse.urljoin(self.base_url, f"svg/{base64_string}")
    req = PreparedRequest()
    req.prepare_url(url, self.query_params)

    if req.url is None:
        raise ValueError("Failed to generate the URL")

    return req.url