Merge branch 'feature/mermaid_suffixes' into 'mgx_ops'

feat: mermaid_to_file + suffixes

See merge request pub/MetaGPT!328

metagpt/utils/__init__.py

@@ -19,6 +19,7 @@ __all__ = [
     "read_docx",
     "Singleton",
     "TOKEN_COSTS",
+    "new_transaction_id",
     "count_message_tokens",
     "count_string_tokens",
 ]

metagpt/utils/common.py

@@ -26,6 +26,7 @@ import re
 import sys
 import time
 import traceback
+import uuid
 from asyncio import iscoroutinefunction
 from datetime import datetime
 from functools import partial
@@ -1089,6 +1090,19 @@ def tool2name(cls, methods: List[str], entry) -> Dict[str, Any]:
     return mappings
 
 
+def new_transaction_id(postfix_len=8) -> str:
+    """
+    Generates a new unique transaction ID based on current timestamp and a random UUID.
+
+    Args:
+        postfix_len (int): Length of the random UUID postfix to include in the transaction ID. Default is 8.
+
+    Returns:
+        str: A unique transaction ID composed of timestamp and a random UUID.
+    """
+    return datetime.now().strftime("%Y%m%d%H%M%ST") + uuid.uuid4().hex[0:postfix_len]
+
+
 def log_time(method):
     """A time-consuming decorator for printing execution duration."""
 

metagpt/utils/mermaid.py

@@ -8,6 +8,7 @@
 import asyncio
 import os
 from pathlib import Path
+from typing import List, Optional
 
 from metagpt.config2 import Config
 from metagpt.logs import logger
@@ -15,16 +16,29 @@ from metagpt.utils.common import awrite, check_cmd_exists
 
 
 async def mermaid_to_file(
-    engine, mermaid_code, output_file_without_suffix, width=2048, height=2048, config=None
+    engine,
+    mermaid_code,
+    output_file_without_suffix,
+    width=2048,
+    height=2048,
+    config=None,
+    suffixes: Optional[List[str]] = None,
 ) -> int:
-    """suffix: png/svg/pdf
+    """Convert Mermaid code to various file formats.
 
-    :param mermaid_code: mermaid code
-    :param output_file_without_suffix: output filename
-    :param width:
-    :param height:
-    :return: 0 if succeed, -1 if failed
+    Args:
+        engine (str): The engine to use for conversion. Supported engines are "nodejs", "playwright", "pyppeteer", "ink", and "none".
+        mermaid_code (str): The Mermaid code to be converted.
+        output_file_without_suffix (str): The output file name without the suffix.
+        width (int, optional): The width of the output image. Defaults to 2048.
+        height (int, optional): The height of the output image. Defaults to 2048.
+        config (Optional[Config], optional): The configuration to use for the conversion. Defaults to None, which uses the default configuration.
+        suffixes (Optional[List[str]], optional): The file suffixes to generate. Supports "png", "pdf", and "svg". Defaults to ["png"].
+
+    Returns:
+        int: 0 if the conversion is successful, -1 if the conversion fails.
     """
+    suffixes = suffixes or ["png"]
     # Write the Mermaid code to a temporary file
     config = config if config else Config.default()
     dir_name = os.path.dirname(output_file_without_suffix)
@@ -41,7 +55,7 @@ async def mermaid_to_file(
             )
             return -1
 
-        for suffix in ["pdf", "svg", "png"]:
+        for suffix in suffixes:
             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}..")
@@ -75,15 +89,15 @@ async def mermaid_to_file(
         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)
+            return await mermaid_to_file(mermaid_code, output_file_without_suffix, width, height, suffixes=suffixes)
         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)
+            return await mermaid_to_file(mermaid_code, output_file_without_suffix, width, height, suffixes=suffixes)
         elif engine == "ink":
             from metagpt.utils.mmdc_ink import mermaid_to_file
 
-            return await mermaid_to_file(mermaid_code, output_file_without_suffix)
+            return await mermaid_to_file(mermaid_code, output_file_without_suffix, suffixes=suffixes)
         elif engine == "none":
             return 0
         else:

metagpt/utils/mmdc_ink.py

@@ -6,21 +6,29 @@
 @File    : mermaid.py
 """
 import base64
+from typing import List, Optional
 
 from aiohttp import ClientError, ClientSession
 
 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
+async def mermaid_to_file(mermaid_code, output_file_without_suffix, suffixes: Optional[List[str]] = None):
+    """Convert Mermaid code to various file formats.
+
+    Args:
+        mermaid_code (str): The Mermaid code to be converted.
+        output_file_without_suffix (str): The output file name without the suffix.
+        width (int, optional): The width of the output image. Defaults to 2048.
+        height (int, optional): The height of the output image. Defaults to 2048.
+        suffixes (Optional[List[str]], optional): The file suffixes to generate. Supports "png", "pdf", and "svg". Defaults to ["png"].
+
+    Returns:
+        int: 0 if the conversion is successful, -1 if the conversion fails.
     """
     encoded_string = base64.b64encode(mermaid_code.encode()).decode()
-
-    for suffix in ["svg", "png"]:
+    suffixes = suffixes or ["png"]
+    for suffix in suffixes:
         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}"

metagpt/utils/mmdc_playwright.py

@@ -7,6 +7,7 @@
 """
 
 import os
+from typing import List, Optional
 from urllib.parse import urljoin
 
 from playwright.async_api import async_playwright
@@ -14,20 +15,22 @@ 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.
+async def mermaid_to_file(
+    mermaid_code, output_file_without_suffix, width=2048, height=2048, suffixes: Optional[List[str]] = None
+) -> int:
+    """Convert Mermaid code to various file formats.
 
     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.
+        mermaid_code (str): The Mermaid code to be converted.
+        output_file_without_suffix (str): The output file name without the suffix.
+        width (int, optional): The width of the output image. Defaults to 2048.
+        height (int, optional): The height of the output image. Defaults to 2048.
+        suffixes (Optional[List[str]], optional): The file suffixes to generate. Supports "png", "pdf", and "svg". Defaults to ["png"].
 
     Returns:
-        int: Returns 1 if the conversion and saving were successful, -1 otherwise.
+        int: 0 if the conversion is successful, -1 if the conversion fails.
     """
-    suffixes = ["png", "svg", "pdf"]
+    suffixes = suffixes or ["png"]
     __dirname = os.path.dirname(os.path.abspath(__file__))
 
     async with async_playwright() as p:

metagpt/utils/mmdc_pyppeteer.py

@@ -6,6 +6,7 @@
 @File    : mmdc_pyppeteer.py
 """
 import os
+from typing import List, Optional
 from urllib.parse import urljoin
 
 from pyppeteer import launch
@@ -14,21 +15,24 @@ from metagpt.config2 import Config
 from metagpt.logs import logger
 
 
-async def mermaid_to_file(mermaid_code, output_file_without_suffix, width=2048, height=2048, config=None) -> int:
-    """
-    Converts the given Mermaid code to various output formats and saves them to files.
+async def mermaid_to_file(
+    mermaid_code, output_file_without_suffix, width=2048, height=2048, config=None, suffixes: Optional[List[str]] = None
+) -> int:
+    """Convert Mermaid code to various file formats.
 
     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.
+        mermaid_code (str): The Mermaid code to be converted.
+        output_file_without_suffix (str): The output file name without the suffix.
+        width (int, optional): The width of the output image. Defaults to 2048.
+        height (int, optional): The height of the output image. Defaults to 2048.
+        config (Optional[Config], optional): The configuration to use for the conversion. Defaults to None, which uses the default configuration.
+        suffixes (Optional[List[str]], optional): The file suffixes to generate. Supports "png", "pdf", and "svg". Defaults to ["png"].
 
     Returns:
-        int: Returns 1 if the conversion and saving were successful, -1 otherwise.
+        int: 0 if the conversion is successful, -1 if the conversion fails.
     """
     config = config if config else Config.default()
-    suffixes = ["png", "svg", "pdf"]
+    suffixes = suffixes or ["png"]
     __dirname = os.path.dirname(os.path.abspath(__file__))
 
     if config.mermaid.pyppeteer_path:

tests/metagpt/utils/test_mermaid.py

@@ -8,28 +8,32 @@
 
 import pytest
 
-from metagpt.utils.common import check_cmd_exists
+from metagpt.const import DEFAULT_WORKSPACE_ROOT
+from metagpt.utils.common import check_cmd_exists, new_transaction_id
 from metagpt.utils.mermaid import MMC1, mermaid_to_file
 
 
 @pytest.mark.asyncio
-@pytest.mark.parametrize("engine", ["nodejs", "ink"])  # TODO: playwright and pyppeteer
-async def test_mermaid(engine, context, mermaid_mocker):
+@pytest.mark.parametrize(
+    ("engine", "suffixes"), [("nodejs", None), ("nodejs", ["png", "svg", "pdf"]), ("ink", None)]
+)  # TODO: playwright and pyppeteer
+async def test_mermaid(engine, suffixes, context, mermaid_mocker):
     # nodejs prerequisites: npm install -g @mermaid-js/mermaid-cli
     # ink prerequisites: connected to internet
     # playwright prerequisites: playwright install --with-deps chromium
     assert check_cmd_exists("npm") == 0
 
-    save_to = context.git_repo.workdir / f"{engine}/1"
-    await mermaid_to_file(engine, MMC1, save_to)
+    save_to = DEFAULT_WORKSPACE_ROOT / f"{new_transaction_id()}/{engine}/1"
+    await mermaid_to_file(engine, MMC1, save_to, suffixes=suffixes)
 
     # ink does not support pdf
+    exts = ["." + i for i in suffixes] if suffixes else [".png"]
     if engine == "ink":
-        for ext in [".svg", ".png"]:
+        for ext in exts:
             assert save_to.with_suffix(ext).exists()
             save_to.with_suffix(ext).unlink(missing_ok=True)
     else:
-        for ext in [".pdf", ".svg", ".png"]:
+        for ext in exts:
             assert save_to.with_suffix(ext).exists()
             save_to.with_suffix(ext).unlink(missing_ok=True)