MetaGPT

1
#!/usr/bin/env python
2
# -*- coding: utf-8 -*-
3
"""
4
@Time    : 2023/7/4 10:53
5
@Author  : alexanderwu alitrack
6
@File    : mermaid.py
7
"""
8
import asyncio
9
import os
10
from pathlib import Path
11

12
from metagpt.config2 import config
13
from metagpt.logs import logger
14
from metagpt.utils.common import awrite, check_cmd_exists
15

16

17
async def mermaid_to_file(engine, mermaid_code, output_file_without_suffix, width=2048, height=2048) -> int:
18
    """suffix: png/svg/pdf
19

20
    :param mermaid_code: mermaid code
21
    :param output_file_without_suffix: output filename
22
    :param width:
23
    :param height:
24
    :return: 0 if succeed, -1 if failed
25
    """
26
    # Write the Mermaid code to a temporary file
27
    dir_name = os.path.dirname(output_file_without_suffix)
28
    if dir_name and not os.path.exists(dir_name):
29
        os.makedirs(dir_name)
30
    tmp = Path(f"{output_file_without_suffix}.mmd")
31
    await awrite(filename=tmp, data=mermaid_code)
32

33
    if engine == "nodejs":
34
        if check_cmd_exists(config.mermaid.path) != 0:
35
            logger.warning(
36
                "RUN `npm install -g @mermaid-js/mermaid-cli` to install mmdc,"
37
                "or consider changing engine to `playwright`, `pyppeteer`, or `ink`."
38
            )
39
            return -1
40

41
        for suffix in ["pdf", "svg", "png"]:
42
            output_file = f"{output_file_without_suffix}.{suffix}"
43
            # Call the `mmdc` command to convert the Mermaid code to a PNG
44
            logger.info(f"Generating {output_file}..")
45

46
            if config.mermaid.puppeteer_config:
47
                commands = [
48
                    config.mermaid.path,
49
                    "-p",
50
                    config.mermaid.puppeteer_config,
51
                    "-i",
52
                    str(tmp),
53
                    "-o",
54
                    output_file,
55
                    "-w",
56
                    str(width),
57
                    "-H",
58
                    str(height),
59
                ]
60
            else:
61
                commands = [config.mermaid.path, "-i", str(tmp), "-o", output_file, "-w", str(width), "-H", str(height)]
62
            process = await asyncio.create_subprocess_shell(
63
                " ".join(commands), stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE
64
            )
65

66
            stdout, stderr = await process.communicate()
67
            if stdout:
68
                logger.info(stdout.decode())
69
            if stderr:
70
                logger.warning(stderr.decode())
71
    else:
72
        if engine == "playwright":
73
            from metagpt.utils.mmdc_playwright import mermaid_to_file
74

75
            return await mermaid_to_file(mermaid_code, output_file_without_suffix, width, height)
76
        elif engine == "pyppeteer":
77
            from metagpt.utils.mmdc_pyppeteer import mermaid_to_file
78

79
            return await mermaid_to_file(mermaid_code, output_file_without_suffix, width, height)
80
        elif engine == "ink":
81
            from metagpt.utils.mmdc_ink import mermaid_to_file
82

83
            return await mermaid_to_file(mermaid_code, output_file_without_suffix)
84
        elif engine == "none":
85
            return 0
86
        else:
87
            logger.warning(f"Unsupported mermaid engine: {engine}")
88
    return 0
89

90

91
MMC1 = """
92
classDiagram
93
    class Main {
94
        -SearchEngine search_engine
95
        +main() str
96
    }
97
    class SearchEngine {
98
        -Index index
99
        -Ranking ranking
100
        -Summary summary
101
        +search(query: str) str
102
    }
103
    class Index {
104
        -KnowledgeBase knowledge_base
105
        +create_index(data: dict)
106
        +query_index(query: str) list
107
    }
108
    class Ranking {
109
        +rank_results(results: list) list
110
    }
111
    class Summary {
112
        +summarize_results(results: list) str
113
    }
114
    class KnowledgeBase {
115
        +update(data: dict)
116
        +fetch_data(query: str) dict
117
    }
118
    Main --> SearchEngine
119
    SearchEngine --> Index
120
    SearchEngine --> Ranking
121
    SearchEngine --> Summary
122
    Index --> KnowledgeBase
123
"""
124

125
MMC2 = """
126
sequenceDiagram
127
    participant M as Main
128
    participant SE as SearchEngine
129
    participant I as Index
130
    participant R as Ranking
131
    participant S as Summary
132
    participant KB as KnowledgeBase
133
    M->>SE: search(query)
134
    SE->>I: query_index(query)
135
    I->>KB: fetch_data(query)
136
    KB-->>I: return data
137
    I-->>SE: return results
138
    SE->>R: rank_results(results)
139
    R-->>SE: return ranked_results
140
    SE->>S: summarize_results(ranked_results)
141
    S-->>SE: return summary
142
    SE-->>M: return summary
143
"""
144