diff --git a/metagpt/utils/__init__.py b/metagpt/utils/__init__.py index f13175cf8..26042eb0e 100644 --- a/metagpt/utils/__init__.py +++ b/metagpt/utils/__init__.py @@ -19,6 +19,7 @@ __all__ = [ "read_docx", "Singleton", "TOKEN_COSTS", + "new_transaction_id", "count_message_tokens", "count_string_tokens", ] diff --git a/metagpt/utils/common.py b/metagpt/utils/common.py index 6ca9b37a2..2b2a209be 100644 --- a/metagpt/utils/common.py +++ b/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.""" diff --git a/metagpt/utils/mermaid.py b/metagpt/utils/mermaid.py index d87ae4f83..64996717e 100644 --- a/metagpt/utils/mermaid.py +++ b/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: diff --git a/metagpt/utils/mmdc_ink.py b/metagpt/utils/mmdc_ink.py index d594adb30..15d6d6083 100644 --- a/metagpt/utils/mmdc_ink.py +++ b/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}" diff --git a/metagpt/utils/mmdc_playwright.py b/metagpt/utils/mmdc_playwright.py index 5d455e1c5..cf846a7e9 100644 --- a/metagpt/utils/mmdc_playwright.py +++ b/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: diff --git a/metagpt/utils/mmdc_pyppeteer.py b/metagpt/utils/mmdc_pyppeteer.py index 4e30ee538..36b77b5b2 100644 --- a/metagpt/utils/mmdc_pyppeteer.py +++ b/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: diff --git a/tests/metagpt/utils/test_mermaid.py b/tests/metagpt/utils/test_mermaid.py index 7367463dc..1fbf060fe 100644 --- a/tests/metagpt/utils/test_mermaid.py +++ b/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)