Merge pull request #309 from alitrack/main

support more options for mermaid conversion

README.md

@@ -93,6 +93,71 @@ # Step 3: Clone the repository to your local machine, and install it.
 
 - if `python setup.py install` fails with error `[Errno 13] Permission denied: '/usr/local/lib/python3.11/dist-packages/test-easy-install-13129.write-test'`, try instead running `python setup.py install --user`
 
+- To convert Mermaid charts to SVG, PNG, and PDF formats. In addition to the Node.js version of Mermaid-CLI, you now have the option to use Python version Playwright, pyppeteer or mermaid.ink for this task.
+
+  - Playwright
+    - **Install Playwright**
+
+    ```bash
+    pip install playwright
+    ```
+
+    - **Install the Required Browsers**
+
+    to support PDF conversion, please install Chrominum.
+
+    ```bash
+    playwright install --with-deps chromium
+    ```
+
+    - **modify `config.yaml`**
+
+    uncomment MERMAID_ENGINE from config.yaml and change it to `playwright`
+
+    ```yaml
+    MERMAID_ENGINE: playwright
+    ```
+
+  - pyppeteer
+    - **Install pyppeteer**
+
+    ```bash
+    pip install pyppeteer
+    ```
+
+    - **Use your own Browsers**
+
+    pyppeteer alow you use installed browsers,  please set the following envirment
+    
+    ```bash
+    export PUPPETEER_EXECUTABLE_PATH = /path/to/your/chromium or edge or chrome
+    ```
+
+    please do not use this command to install browser, it is too old
+
+    ```bash
+    pyppeteer-install
+    ```
+
+    - **modify `config.yaml`**
+
+    uncomment MERMAID_ENGINE from config.yaml and change it to `pyppeteer`
+
+    ```yaml
+    MERMAID_ENGINE: pyppeteer
+    ```
+
+  - mermaid.ink
+    - **modify `config.yaml`**
+
+    uncomment MERMAID_ENGINE from config.yaml and change it to `ink`
+
+    ```yaml
+    MERMAID_ENGINE: ink
+    ```  
+
+    Note: this method does not support pdf export.
+
 ### Installation by Docker
 
 ```bash

config/config.yaml

@@ -76,3 +76,10 @@ SD_T2I_API: "/sdapi/v1/txt2img"
 ### for Research
 MODEL_FOR_RESEARCHER_SUMMARY: gpt-3.5-turbo
 MODEL_FOR_RESEARCHER_REPORT: gpt-3.5-turbo-16k
+
+### choose the engine for mermaid conversion, 
+# default is nodejs, you can change it to playwright,pyppeteer or ink
+# MERMAID_ENGINE: nodejs
+
+### browser path for pyppeteer engine, support Chrome, Chromium,MS Edge
+#PYPPETEER_EXECUTABLE_PATH: "/usr/bin/google-chrome-stable"

metagpt/actions/design_api.py

@@ -103,23 +103,23 @@ class WriteDesign(Action):
             pass  # Folder does not exist, but we don't care
         workspace.mkdir(parents=True, exist_ok=True)
 
-    def _save_prd(self, docs_path, resources_path, prd):
+    async  def _save_prd(self, docs_path, resources_path, prd):
         prd_file = docs_path / 'prd.md'
         quadrant_chart = CodeParser.parse_code(block="Competitive Quadrant Chart", text=prd)
-        mermaid_to_file(quadrant_chart, resources_path / 'competitive_analysis')
+        await mermaid_to_file(quadrant_chart, resources_path / 'competitive_analysis')
         logger.info(f"Saving PRD to {prd_file}")
         prd_file.write_text(prd)
 
-    def _save_system_design(self, docs_path, resources_path, content):
+    async  def _save_system_design(self, docs_path, resources_path, content):
         data_api_design = CodeParser.parse_code(block="Data structures and interface definitions", text=content)
         seq_flow = CodeParser.parse_code(block="Program call flow", text=content)
-        mermaid_to_file(data_api_design, resources_path / 'data_api_design')
-        mermaid_to_file(seq_flow, resources_path / 'seq_flow')
+        await mermaid_to_file(data_api_design, resources_path / 'data_api_design')
+        await mermaid_to_file(seq_flow, resources_path / 'seq_flow')
         system_design_file = docs_path / 'system_design.md'
         logger.info(f"Saving System Designs to {system_design_file}")
         system_design_file.write_text(content)
 
-    def _save(self, context, system_design):
+    async  def _save(self, context, system_design):
         if isinstance(system_design, ActionOutput):
             content = system_design.content
             ws_name = CodeParser.parse_str(block="Python package name", text=content)
@@ -132,13 +132,13 @@ class WriteDesign(Action):
         resources_path = workspace / 'resources'
         docs_path.mkdir(parents=True, exist_ok=True)
         resources_path.mkdir(parents=True, exist_ok=True)
-        self._save_prd(docs_path, resources_path, context[-1].content)
-        self._save_system_design(docs_path, resources_path, content)
+        await self._save_prd(docs_path, resources_path, context[-1].content)
+        await self._save_system_design(docs_path, resources_path, content)
 
     async def run(self, context):
         prompt = PROMPT_TEMPLATE.format(context=context, format_example=FORMAT_EXAMPLE)
         # system_design = await self._aask(prompt)
         system_design = await self._aask_v1(prompt, "system_design", OUTPUT_MAPPING)
-        self._save(context, system_design)
+        await self._save(context, system_design)
         return system_design
     

metagpt/config.py

@@ -83,6 +83,8 @@ class Config(metaclass=Singleton):
         self.calc_usage = self._get("CALC_USAGE", True)
         self.model_for_researcher_summary = self._get("MODEL_FOR_RESEARCHER_SUMMARY")
         self.model_for_researcher_report = self._get("MODEL_FOR_RESEARCHER_REPORT")
+        self.mermaid_engine = self._get("MERMAID_ENGINE", 'nodejs')
+        self.pyppeteer_executable_path = self._get("PYPPETEER_EXECUTABLE_PATH", '')
 
     def _init_with_config_files_and_env(self, configs: dict, yaml_file):
         """Load from config/key.yaml, config/config.yaml, and env in decreasing order of priority"""

metagpt/utils/common.py

@@ -191,7 +191,8 @@ class CodeParser:
         else:
             logger.error(f"{pattern} not match following text:")
             logger.error(text)
-            raise Exception
+            # raise Exception
+            return ""
         return code
 
     @classmethod

metagpt/utils/mermaid.py

@@ -2,9 +2,10 @@
 # -*- coding: utf-8 -*-
 """
 @Time    : 2023/7/4 10:53
-@Author  : alexanderwu
+@Author  : alexanderwu alitrack
 @File    : mermaid.py
 """
+import asyncio
 import subprocess
 from pathlib import Path
 
@@ -12,48 +13,76 @@ from metagpt.config import CONFIG
 from metagpt.const import PROJECT_ROOT
 from metagpt.logs import logger
 from metagpt.utils.common import check_cmd_exists
+import os
+import sys
 
-
-def mermaid_to_file(mermaid_code, output_file_without_suffix, width=2048, height=2048) -> int:
+async def mermaid_to_file(mermaid_code, output_file_without_suffix, width=2048, height=2048) -> int:
     """suffix: png/svg/pdf
 
     :param mermaid_code: mermaid code
     :param output_file_without_suffix: output filename
     :param width:
     :param height:
-    :return: 0 if succed, -1 if failed
+    :return: 0 if succeed, -1 if failed
     """
     # Write the Mermaid code to a temporary file
+    dir_name = os.path.dirname(output_file_without_suffix)
+    if dir_name and not os.path.exists(dir_name):
+        os.makedirs(dir_name)    
     tmp = Path(f"{output_file_without_suffix}.mmd")
     tmp.write_text(mermaid_code, encoding="utf-8")
+    
+    engine = CONFIG.mermaid_engine.lower()
+    if engine == "nodejs":
+        if check_cmd_exists("mmdc") != 0:
+            logger.warning("RUN `npm install -g @mermaid-js/mermaid-cli` to install mmdc")
+            return -1
+        
+        for suffix in ["pdf", "svg", "png"]:
+            output_file = f"{output_file_without_suffix}.{suffix}"
+            # Call the `mmdc` command to convert the Mermaid code to a PNG
+            logger.info(f"Generating {output_file}..")
 
-    if check_cmd_exists("mmdc") != 0:
-        logger.warning("RUN `npm install -g @mermaid-js/mermaid-cli` to install mmdc")
-        return -1
-
-    for suffix in ["pdf", "svg", "png"]:
-        output_file = f"{output_file_without_suffix}.{suffix}"
-        # Call the `mmdc` command to convert the Mermaid code to a PNG
-        logger.info(f"Generating {output_file}..")
-
-        if CONFIG.puppeteer_config:
-            subprocess.run(
-                [
-                    CONFIG.mmdc,
-                    "-p",
-                    CONFIG.puppeteer_config,
-                    "-i",
-                    str(tmp),
-                    "-o",
-                    output_file,
-                    "-w",
-                    str(width),
-                    "-H",
-                    str(height),
-                ]
+            if CONFIG.puppeteer_config:
+                commands =[
+                        CONFIG.mmdc,
+                        "-p",
+                        CONFIG.puppeteer_config,
+                        "-i",
+                        str(tmp),
+                        "-o",
+                        output_file,
+                        "-w",
+                        str(width),
+                        "-H",
+                        str(height),
+                    ]
+            else:
+                commands =[CONFIG.mmdc, "-i", str(tmp), "-o", output_file, "-w", str(width), "-H", str(height)]
+            process = await asyncio.create_subprocess_exec(
+                *commands,
+                stdout=asyncio.subprocess.PIPE,
+                stderr=asyncio.subprocess.PIPE
             )
+
+            stdout, stderr = await process.communicate()
+            if stdout:
+                logger.info(stdout.decode())
+            if stderr:
+                logger.error(stderr.decode())
+    else:
+
+        if engine =='playwright':
+            from metagpt.utils.mmdc_playwright import mermaid_to_file
+            return await mermaid_to_file(mermaid_code, output_file_without_suffix, width, height)
+        elif engine =='pyppeteer':
+            from metagpt.utils.mmdc_pyppeteer import mermaid_to_file
+            return await mermaid_to_file(mermaid_code, output_file_without_suffix, width, height)
+        elif engine =='ink':
+            from metagpt.utils.mmdc_ink import mermaid_to_file
+            return await mermaid_to_file(mermaid_code, output_file_without_suffix)
         else:
-            subprocess.run([CONFIG.mmdc, "-i", str(tmp), "-o", output_file, "-w", str(width), "-H", str(height)])
+            logger.warning(f"Unsupported mermaid engine: {engine}")
     return 0
 
 
@@ -108,7 +137,9 @@ MMC2 = """sequenceDiagram
     SE-->>M: return summary"""
 
 
+
 if __name__ == "__main__":
-    # logger.info(print_members(print_members))
-    mermaid_to_file(MMC1, PROJECT_ROOT / "tmp/1.png")
-    mermaid_to_file(MMC2, PROJECT_ROOT / "tmp/2.png")
+    loop = asyncio.new_event_loop()
+    result  = loop.run_until_complete(mermaid_to_file(MMC1, PROJECT_ROOT / f"{CONFIG.mermaid_engine}/1"))
+    result  = loop.run_until_complete(mermaid_to_file(MMC2, PROJECT_ROOT / f"{CONFIG.mermaid_engine}/1"))
+    loop.close()

metagpt/utils/mmdc_ink.py

@@ -0,0 +1,41 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@Time    : 2023/9/4 16:12
+@Author  : alitrack
+@File    : mermaid.py
+"""
+import base64
+import os
+
+from aiohttp import ClientSession,ClientError
+from metagpt.logs import logger
+
+
+async def mermaid_to_file(mermaid_code, output_file_without_suffix):
+    """suffix: png/svg
+    :param mermaid_code: mermaid code
+    :param output_file_without_suffix: output filename without suffix
+    :return: 0 if succeed, -1 if failed
+    """
+    encoded_string = base64.b64encode(mermaid_code.encode()).decode()
+
+    for suffix in ["svg", "png"]:
+        output_file = f"{output_file_without_suffix}.{suffix}"
+        path_type = "svg" if suffix == "svg" else "img"
+        url = f"https://mermaid.ink/{path_type}/{encoded_string}"
+        async with ClientSession() as session:
+            try:
+                async with session.get(url) as response:
+                    if response.status == 200:
+                        text = await response.content.read()
+                        with open(output_file, 'wb') as f:
+                            f.write(text)
+                        logger.info(f"Generating {output_file}..")
+                    else:
+                        logger.error(f"Failed to generate {output_file}")
+                        return -1
+            except ClientError as e:
+                logger.error(f"network error: {e}")
+                return -1
+    return 0

metagpt/utils/mmdc_playwright.py

@@ -0,0 +1,111 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@Time    : 2023/9/4 16:12
+@Author  : Steven Lee
+@File    : mmdc_playwright.py
+"""
+
+import os
+from urllib.parse import urljoin
+from playwright.async_api import async_playwright
+from metagpt.logs import logger
+
+async def mermaid_to_file(mermaid_code, output_file_without_suffix, width=2048, height=2048)-> int:
+    """
+    Converts the given Mermaid code to various output formats and saves them to files.
+
+    Args:
+        mermaid_code (str): The Mermaid code to convert.
+        output_file_without_suffix (str): The output file name without the file extension.
+        width (int, optional): The width of the output image in pixels. Defaults to 2048.
+        height (int, optional): The height of the output image in pixels. Defaults to 2048.
+
+    Returns:
+        int: Returns 1 if the conversion and saving were successful, -1 otherwise.
+    """
+    suffixes=['png', 'svg', 'pdf']
+    __dirname = os.path.dirname(os.path.abspath(__file__))
+
+    async with async_playwright() as p:
+        browser = await p.chromium.launch()
+        device_scale_factor = 1.0
+        context = await browser.new_context(
+                viewport={'width': width, 'height': height},
+                device_scale_factor=device_scale_factor,
+            )
+        page = await context.new_page()
+
+        async def console_message(msg):
+            logger.info(msg.text)
+        page.on('console', console_message)
+
+        try:
+            await page.set_viewport_size({'width': width, 'height': height})
+
+            mermaid_html_path = os.path.abspath(
+                os.path.join(__dirname, 'index.html'))
+            mermaid_html_url = urljoin('file:', mermaid_html_path)
+            await page.goto(mermaid_html_url)
+            await page.wait_for_load_state("networkidle")
+
+            await page.wait_for_selector("div#container", state="attached")
+            mermaid_config = {}
+            background_color = "#ffffff"
+            my_css = ""
+            await page.evaluate(f'document.body.style.background = "{background_color}";')
+
+            metadata = await page.evaluate('''async ([definition, mermaidConfig, myCSS, backgroundColor]) => {
+                const { mermaid, zenuml } = globalThis;
+                await mermaid.registerExternalDiagrams([zenuml]);
+                mermaid.initialize({ startOnLoad: false, ...mermaidConfig });
+                const { svg } = await mermaid.render('my-svg', definition, document.getElementById('container'));
+                document.getElementById('container').innerHTML = svg;
+                const svgElement = document.querySelector('svg');
+                svgElement.style.backgroundColor = backgroundColor;
+
+                if (myCSS) {
+                    const style = document.createElementNS('http://www.w3.org/2000/svg', 'style');
+                    style.appendChild(document.createTextNode(myCSS));
+                    svgElement.appendChild(style);
+                }
+
+            }''', [mermaid_code, mermaid_config, my_css, background_color])
+
+            if 'svg' in suffixes :
+                svg_xml = await page.evaluate('''() => {
+                    const svg = document.querySelector('svg');
+                    const xmlSerializer = new XMLSerializer();
+                    return xmlSerializer.serializeToString(svg);
+                }''')
+                logger.info(f"Generating {output_file_without_suffix}.svg..")
+                with open(f'{output_file_without_suffix}.svg', 'wb') as f:
+                    f.write(svg_xml.encode('utf-8'))
+
+            if  'png' in suffixes:
+                clip = await page.evaluate('''() => {
+                    const svg = document.querySelector('svg');
+                    const rect = svg.getBoundingClientRect();
+                    return {
+                        x: Math.floor(rect.left),
+                        y: Math.floor(rect.top),
+                        width: Math.ceil(rect.width),
+                        height: Math.ceil(rect.height)
+                    };
+                }''')
+                await page.set_viewport_size({'width': clip['x'] + clip['width'], 'height': clip['y'] + clip['height']})
+                screenshot = await page.screenshot(clip=clip, omit_background=True, scale='device')
+                logger.info(f"Generating {output_file_without_suffix}.png..")
+                with open(f'{output_file_without_suffix}.png', 'wb') as f:
+                    f.write(screenshot)
+            if 'pdf' in suffixes:
+                pdf_data = await page.pdf(scale=device_scale_factor)
+                logger.info(f"Generating {output_file_without_suffix}.pdf..")
+                with open(f'{output_file_without_suffix}.pdf', 'wb') as f:
+                    f.write(pdf_data)
+            return 0
+        except Exception as e:
+            logger.error(e)
+            return -1
+        finally:
+            await browser.close()

metagpt/utils/mmdc_pyppeteer.py

@@ -0,0 +1,113 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@Time    : 2023/9/4 16:12
+@Author  : alitrack
+@File    : mmdc_pyppeteer.py
+"""
+import os
+from urllib.parse import urljoin
+from pyppeteer import launch
+from metagpt.logs import logger
+from metagpt.config import CONFIG
+
+async def mermaid_to_file(mermaid_code, output_file_without_suffix, width=2048, height=2048)-> int:
+    """
+    Converts the given Mermaid code to various output formats and saves them to files.
+
+    Args:
+        mermaid_code (str): The Mermaid code to convert.
+        output_file_without_suffix (str): The output file name without the file extension.
+        width (int, optional): The width of the output image in pixels. Defaults to 2048.
+        height (int, optional): The height of the output image in pixels. Defaults to 2048.
+
+    Returns:
+        int: Returns 1 if the conversion and saving were successful, -1 otherwise.
+    """
+    suffixes = ['png', 'svg', 'pdf']   
+    __dirname = os.path.dirname(os.path.abspath(__file__))
+
+    
+    if CONFIG.pyppeteer_executable_path:
+        browser = await launch(headless=True,
+                            executablePath=CONFIG.pyppeteer_executable_path,
+                            args=['--disable-extensions',"--no-sandbox"] 
+                            )
+    else:
+        logger.error("Please set the environment variable:PYPPETEER_EXECUTABLE_PATH.")
+        return -1
+    page = await browser.newPage()
+    device_scale_factor = 1.0
+
+    async def console_message(msg):
+        logger.info(msg.text)
+    page.on('console', console_message)
+
+    try:
+        await page.setViewport(viewport={'width': width, 'height': height, 'deviceScaleFactor': device_scale_factor})
+
+        mermaid_html_path = os.path.abspath(
+            os.path.join(__dirname, 'index.html'))
+        mermaid_html_url = urljoin('file:', mermaid_html_path)
+        await page.goto(mermaid_html_url)
+
+        await page.querySelector("div#container")
+        mermaid_config = {}
+        background_color = "#ffffff"
+        my_css = ""
+        await page.evaluate(f'document.body.style.background = "{background_color}";')
+
+        metadata = await page.evaluate('''async ([definition, mermaidConfig, myCSS, backgroundColor]) => {
+            const { mermaid, zenuml } = globalThis;
+            await mermaid.registerExternalDiagrams([zenuml]);
+            mermaid.initialize({ startOnLoad: false, ...mermaidConfig });
+            const { svg } = await mermaid.render('my-svg', definition, document.getElementById('container'));
+            document.getElementById('container').innerHTML = svg;
+            const svgElement = document.querySelector('svg');
+            svgElement.style.backgroundColor = backgroundColor;
+
+            if (myCSS) {
+                const style = document.createElementNS('http://www.w3.org/2000/svg', 'style');
+                style.appendChild(document.createTextNode(myCSS));
+                svgElement.appendChild(style);
+            }
+        }''', [mermaid_code, mermaid_config, my_css, background_color])
+
+        if 'svg' in suffixes :
+            svg_xml = await page.evaluate('''() => {
+                const svg = document.querySelector('svg');
+                const xmlSerializer = new XMLSerializer();
+                return xmlSerializer.serializeToString(svg);
+            }''')
+            logger.info(f"Generating {output_file_without_suffix}.svg..")
+            with open(f'{output_file_without_suffix}.svg', 'wb') as f:
+                f.write(svg_xml.encode('utf-8'))
+
+        if  'png' in suffixes:
+            clip = await page.evaluate('''() => {
+                const svg = document.querySelector('svg');
+                const rect = svg.getBoundingClientRect();
+                return {
+                    x: Math.floor(rect.left),
+                    y: Math.floor(rect.top),
+                    width: Math.ceil(rect.width),
+                    height: Math.ceil(rect.height)
+                };
+            }''')
+            await page.setViewport({'width': clip['x'] + clip['width'], 'height': clip['y'] + clip['height'], 'deviceScaleFactor': device_scale_factor})
+            screenshot = await page.screenshot(clip=clip, omit_background=True, scale='device')
+            logger.info(f"Generating {output_file_without_suffix}.png..")
+            with open(f'{output_file_without_suffix}.png', 'wb') as f:
+                f.write(screenshot)
+        if 'pdf' in suffixes:
+            pdf_data = await page.pdf(scale=device_scale_factor)
+            logger.info(f"Generating {output_file_without_suffix}.pdf..")
+            with open(f'{output_file_without_suffix}.pdf', 'wb') as f:
+                f.write(pdf_data)
+        return 0
+    except Exception as e:
+        logger.error(e)
+        return -1
+    finally:
+        await browser.close()
+

setup.py

@@ -47,6 +47,7 @@ setup(
         "selenium": ["selenium>4", "webdriver_manager", "beautifulsoup4"],
         "search-google": ["google-api-python-client==2.94.0"],
         "search-ddg": ["duckduckgo-search==3.8.5"],
+        "pyppeteer": ["pyppeteer>=1.0.2"],
     },
     cmdclass={
         "install_mermaid": InstallMermaidCLI,