MetaGPT

1
#!/usr/bin/env python
2
# -*- coding: utf-8 -*-
3
"""
4
@Time    : 2023/9/4 16:12
5
@Author  : alitrack
6
@File    : mermaid.py
7
"""
8
import base64
9

10
from aiohttp import ClientError, ClientSession
11

12
from metagpt.logs import logger
13

14

15
async def mermaid_to_file(mermaid_code, output_file_without_suffix):
16
    """suffix: png/svg
17
    :param mermaid_code: mermaid code
18
    :param output_file_without_suffix: output filename without suffix
19
    :return: 0 if succeed, -1 if failed
20
    """
21
    encoded_string = base64.b64encode(mermaid_code.encode()).decode()
22

23
    for suffix in ["svg", "png"]:
24
        output_file = f"{output_file_without_suffix}.{suffix}"
25
        path_type = "svg" if suffix == "svg" else "img"
26
        url = f"https://mermaid.ink/{path_type}/{encoded_string}"
27
        async with ClientSession() as session:
28
            try:
29
                async with session.get(url) as response:
30
                    if response.status == 200:
31
                        text = await response.content.read()
32
                        with open(output_file, "wb") as f:
33
                            f.write(text)
34
                        logger.info(f"Generating {output_file}..")
35
                    else:
36
                        logger.error(f"Failed to generate {output_file}")
37
                        return -1
38
            except ClientError as e:
39
                logger.error(f"network error: {e}")
40
                return -1
41
    return 0
42