diff --git a/metagpt/tools/libs/gpt_v_generator.py b/metagpt/tools/libs/gpt_v_generator.py index 3b17fc596..0e9f34770 100644 --- a/metagpt/tools/libs/gpt_v_generator.py +++ b/metagpt/tools/libs/gpt_v_generator.py @@ -5,13 +5,13 @@ @Author : mannaandpoem @File : gpt_v_generator.py """ -import os +import re from pathlib import Path from metagpt.const import DEFAULT_WORKSPACE_ROOT +from metagpt.logs import logger from metagpt.tools.tool_registry import register_tool -from metagpt.tools.tool_type import ToolType -from metagpt.utils.common import encode_image +from metagpt.utils.common import CodeParser, encode_image ANALYZE_LAYOUT_PROMPT = """You are now a UI/UX designer, please generate layout information for this image: @@ -28,11 +28,9 @@ As the design pays tribute to large companies, sometimes it is normal for some c Now, please generate the corresponding webpage code including HTML, CSS and JavaScript:""" -@register_tool( - tool_type=ToolType.IMAGE2WEBPAGE.type_name, include_functions=["__init__", "generate_webpages", "save_webpages"] -) +@register_tool(include_functions=["__init__", "generate_webpages", "save_webpages"]) class GPTvGenerator: - """Class for generating webpages at once. + """Class for generating webpage code from a given webpage screenshot. This class provides methods to generate webpages including all code (HTML, CSS, and JavaScript) based on an image. It utilizes a vision model to analyze the layout from an image and generate webpage codes accordingly. @@ -75,50 +73,34 @@ class GPTvGenerator: return await self.llm.aask(msg=prompt, images=[encode_image(image_path)]) @staticmethod - def save_webpages(image_path: str, webpages: str) -> Path: + def save_webpages(webpages: str, save_folder_name: str = "example") -> Path: """Save webpages including all code (HTML, CSS, and JavaScript) at once. Args: - image_path (str): The path of the image file. webpages (str): The generated webpages content. + save_folder_name (str, optional): The name of the folder to save the webpages. Defaults to 'example'. Returns: Path: The path of the saved webpages. """ # Create a folder called webpages in the workspace directory to store HTML, CSS, and JavaScript files - webpages_path = DEFAULT_WORKSPACE_ROOT / "webpages" / Path(image_path).stem - os.makedirs(webpages_path, exist_ok=True) + webpages_path = DEFAULT_WORKSPACE_ROOT / "webpages" / save_folder_name + logger.info(f"code will be saved at {webpages_path}") + webpages_path.mkdir(parents=True, exist_ok=True) index_path = webpages_path / "index.html" - try: - index = webpages.split("```html")[1].split("```")[0] - style_path = None - if "styles.css" in index: - style_path = webpages_path / "styles.css" - elif "style.css" in index: - style_path = webpages_path / "style.css" - style = webpages.split("```css")[1].split("```")[0] if style_path else "" + index_path.write_text(CodeParser.parse_code(block=None, text=webpages, lang="html")) - js_path = None - if "scripts.js" in index: - js_path = webpages_path / "scripts.js" - elif "script.js" in index: - js_path = webpages_path / "script.js" + extract_and_save_code(folder=webpages_path, text=webpages, pattern="styles?.css", language="css") - js = webpages.split("```javascript")[1].split("```")[0] if js_path else "" - except IndexError: - raise ValueError(f"No html or css or js code found in the result. \nWebpages: {webpages}") - - try: - with open(index_path, "w", encoding="utf-8") as f: - f.write(index) - if style_path: - with open(style_path, "w", encoding="utf-8") as f: - f.write(style) - if js_path: - with open(js_path, "w", encoding="utf-8") as f: - f.write(js) - except FileNotFoundError as e: - raise FileNotFoundError(f"Cannot save the webpages to {str(webpages_path)}") from e + extract_and_save_code(folder=webpages_path, text=webpages, pattern="scripts?.js", language="javascript") return webpages_path + + +def extract_and_save_code(folder, text, pattern, language): + word = re.search(pattern, text) + if word: + path = folder / word.group(0) + code = CodeParser.parse_code(block=None, text=text, lang=language) + path.write_text(code, encoding="utf-8") diff --git a/tests/metagpt/tools/libs/test_gpt_v_generator.py b/tests/metagpt/tools/libs/test_gpt_v_generator.py index 907006765..4a2e68682 100644 --- a/tests/metagpt/tools/libs/test_gpt_v_generator.py +++ b/tests/metagpt/tools/libs/test_gpt_v_generator.py @@ -60,18 +60,24 @@ async def test_generate_webpages(mock_webpage_filename_with_styles_and_scripts, async def test_save_webpages_with_styles_and_scripts(mock_webpage_filename_with_styles_and_scripts, image_path): generator = GPTvGenerator() webpages = await generator.generate_webpages(image_path) - webpages_dir = generator.save_webpages(image_path=image_path, webpages=webpages) + webpages_dir = generator.save_webpages(webpages=webpages, save_folder_name="test_1") logs.logger.info(webpages_dir) assert webpages_dir.exists() + assert (webpages_dir / "index.html").exists() + assert (webpages_dir / "styles.css").exists() + assert (webpages_dir / "scripts.js").exists() @pytest.mark.asyncio async def test_save_webpages_with_style_and_script(mock_webpage_filename_with_style_and_script, image_path): generator = GPTvGenerator() webpages = await generator.generate_webpages(image_path) - webpages_dir = generator.save_webpages(image_path=image_path, webpages=webpages) + webpages_dir = generator.save_webpages(webpages=webpages, save_folder_name="test_2") logs.logger.info(webpages_dir) assert webpages_dir.exists() + assert (webpages_dir / "index.html").exists() + assert (webpages_dir / "style.css").exists() + assert (webpages_dir / "script.js").exists() @pytest.mark.asyncio