From d83d1e105c28b558f9da5e56967834de6b8fe3c6 Mon Sep 17 00:00:00 2001 From: Stitch-z <284618289@qq.com> Date: Thu, 7 Sep 2023 19:55:50 +0800 Subject: [PATCH] update: format directory structure and extract universal file operation class --- metagpt/actions/write_tutorial.py | 67 +++++++++----------- metagpt/prompts/tutorial_assistant.py | 4 +- metagpt/roles/tutorial_assistant.py | 11 ++-- metagpt/utils/file.py | 35 ++++++++++ tests/metagpt/actions/test_write_tutorial.py | 17 +---- tests/metagpt/utils/test_file.py | 27 ++++++++ 6 files changed, 102 insertions(+), 59 deletions(-) create mode 100644 metagpt/utils/file.py create mode 100644 tests/metagpt/utils/test_file.py diff --git a/metagpt/actions/write_tutorial.py b/metagpt/actions/write_tutorial.py index 38a45c4c3..b23fc2ad4 100644 --- a/metagpt/actions/write_tutorial.py +++ b/metagpt/actions/write_tutorial.py @@ -7,13 +7,9 @@ @Describe : Actions of the tutorial assistant, including writing directories and document content. """ import json -from datetime import datetime from typing import Dict -import aiofiles - from metagpt.actions import Action -from metagpt.const import TUTORIAL_PATH from metagpt.logs import logger from metagpt.prompts.tutorial_assistant import DIRECTORY_PROMPT, CONTENT_PROMPT @@ -30,6 +26,33 @@ class WriteDirectory(Action): super().__init__(name, *args, **kwargs) self.language = language + @staticmethod + async def _handle_resp(resp: str) -> Dict: + """Process string results and convert them to JSON format. + + Args: + resp: The directory results returned by gpt. + + Returns: + The parsed dictionary, such as {"title": "xxx", "directory": [{"dir 1": ["sub dir 1", "sub dir 2"]}]}. + + Raises: + Exception: If no matching dictionary section is found. + json.JSONDecodeError: If the dictionary part cannot be parsed as JSON. + """ + start = resp.find('{') + end = resp.rfind('}') + if start != -1 and end != -1 and end > start: + directory_str = resp[start:end + 1] + logger.info(f"Successfully parsed json: {str(directory_str)}") + try: + return json.loads(directory_str) + except json.JSONDecodeError as e: + logger.error(f"Json parsing error: {e}") + raise e + else: + raise Exception("No matching dictionary section found.") + async def run(self, topic: str, *args, **kwargs) -> Dict: """Execute the action to generate a tutorial directory according to the topic. @@ -37,11 +60,11 @@ class WriteDirectory(Action): topic: The tutorial topic. Returns: - the tutorial directory information, such as {"title": "xxx", "directory": [{"dir 1": ["sub dir 1", "sub dir 2"]}]} + the tutorial directory information, including {"title": "xxx", "directory": [{"dir 1": ["sub dir 1", "sub dir 2"]}]}. """ prompt = DIRECTORY_PROMPT.format(topic=topic, language=self.language) - directory = await self._aask(prompt=prompt) - return json.loads(directory) + resp = await self._aask(prompt=prompt) + return await self._handle_resp(resp) class WriteContent(Action): @@ -70,33 +93,3 @@ class WriteContent(Action): prompt = CONTENT_PROMPT.format(topic=topic, language=self.language, directory=self.directory) return await self._aask(prompt=prompt) - -class SaveDocx(Action): - """Action class for saving tutorial docx. - - Args: - name: The name of the action. - """ - - def __init__(self, name: str = "", *args, **kwargs): - super().__init__(name, *args, **kwargs) - - async def run(self, title: str, content: str, *args, **kwargs) -> str: - """Execute the action to save the generated tutorial document to a Markdown file. - - Args: - title: The title of tutorial. - content: The total content of tutorial. - - Returns: - The full filename of tutorial content. - - """ - current_time = datetime.now().strftime("%Y-%m-%d_%H-%M-%S") - pathname = TUTORIAL_PATH / current_time - pathname.mkdir(parents=True, exist_ok=True) - filename = f"{pathname}/{title}.md" - async with aiofiles.open(filename, mode="w", encoding="utf-8") as writer: - await writer.write(content) - logger.info(f"Successfully write docx: {filename}") - return filename \ No newline at end of file diff --git a/metagpt/prompts/tutorial_assistant.py b/metagpt/prompts/tutorial_assistant.py index c9039fd41..d690aad83 100644 --- a/metagpt/prompts/tutorial_assistant.py +++ b/metagpt/prompts/tutorial_assistant.py @@ -15,7 +15,7 @@ We need you to write a technical tutorial with the topic "{topic}". DIRECTORY_PROMPT = COMMON_PROMPT + """ Please provide the specific table of contents for this tutorial, strictly following the following requirements: 1. The output must be strictly in the specified language, {language}. -2. Answer in the dictionary format like {{"title": "xxx", "directory": [{{"dir 1": ["sub dir 1", "sub dir 2"]}}]}}. +2. Answer strictly in the dictionary format like {{"title": "xxx", "directory": [{{"dir 1": ["sub dir 1", "sub dir 2"]}}, {{"dir 2": ["sub dir 3", "sub dir 4"]}}]}}. 3. The directory should be as specific and sufficient as possible, with a primary and secondary directory.The secondary directory is in the array. 4. Do not have extra spaces or line breaks. 5. Each directory title has practical significance. @@ -35,5 +35,5 @@ Strictly limit output according to the following requirements: 2. If there are code examples, they must follow standard syntax specifications, have document annotations, and be displayed in code blocks. 3. The output must be strictly in the specified language, {language}. 4. Do not have redundant output, including concluding remarks. -5. Don't return the topic "{topic}". +5. Strict requirement not to output the topic "{topic}". """ \ No newline at end of file diff --git a/metagpt/roles/tutorial_assistant.py b/metagpt/roles/tutorial_assistant.py index daf4daf40..9a7df4f4d 100644 --- a/metagpt/roles/tutorial_assistant.py +++ b/metagpt/roles/tutorial_assistant.py @@ -6,12 +6,15 @@ @File : tutorial_assistant.py """ +from datetime import datetime from typing import Dict -from metagpt.actions.write_tutorial import WriteDirectory, WriteContent, SaveDocx +from metagpt.actions.write_tutorial import WriteDirectory, WriteContent +from metagpt.const import TUTORIAL_PATH from metagpt.logs import logger from metagpt.roles import Role from metagpt.schema import Message +from metagpt.utils.file import File class TutorialAssistant(Role): @@ -71,7 +74,6 @@ class TutorialAssistant(Role): directory += f"- {key}\n" for second_dir in first_dir[key]: directory += f" - {second_dir}\n" - actions.append(SaveDocx()) self._init_actions(actions) self._rc.todo = None return Message(content=directory) @@ -89,9 +91,6 @@ class TutorialAssistant(Role): resp = await todo.run(topic=self.topic) logger.info(resp) return await self._handle_directory(resp) - elif type(todo) is SaveDocx: - filename = await todo.run(title=self.main_title, content=self.total_content) - return Message(content=filename, role=self.profile) resp = await todo.run(topic=self.topic) logger.info(resp) if self.total_content != "": @@ -110,4 +109,6 @@ class TutorialAssistant(Role): if self._rc.todo is None: break msg = await self._act() + root_path = TUTORIAL_PATH / datetime.now().strftime("%Y-%m-%d_%H-%M-%S") + await File.write(root_path, f"{self.main_title}.md", self.total_content.encode('utf-8')) return msg diff --git a/metagpt/utils/file.py b/metagpt/utils/file.py new file mode 100644 index 000000000..93e1ad6c7 --- /dev/null +++ b/metagpt/utils/file.py @@ -0,0 +1,35 @@ +#!/usr/bin/env python3 +# _*_ coding: utf-8 _*_ +""" +@Time : 2023/9/4 15:40:40 +@Author : Stitch-z +@File : file.py +@Describe : General file operations. +""" +import aiofiles +from pathlib import Path + +from metagpt.logs import logger + + +class File: + """A general util for file operations.""" + + @classmethod + async def write(cls, root_path: Path, filename: str, content: bytes) -> Path: + """Write the file content to the local specified path. + + Args: + root_path: The root path of file, such as "/data". + filename: The name of file, such as "test.txt". + content: The binary content of file. + + Returns: + The full filename of file, such as "/data/test.txt". + """ + root_path.mkdir(parents=True, exist_ok=True) + full_path = root_path / filename + async with aiofiles.open(full_path, mode="wb") as writer: + await writer.write(content) + logger.info(f"Successfully write docx: {full_path}") + return full_path \ No newline at end of file diff --git a/tests/metagpt/actions/test_write_tutorial.py b/tests/metagpt/actions/test_write_tutorial.py index 6460aa08b..683fee082 100644 --- a/tests/metagpt/actions/test_write_tutorial.py +++ b/tests/metagpt/actions/test_write_tutorial.py @@ -7,10 +7,9 @@ """ from typing import Dict -import aiofiles import pytest -from metagpt.actions.write_tutorial import WriteDirectory, WriteContent, SaveDocx +from metagpt.actions.write_tutorial import WriteDirectory, WriteContent @pytest.mark.asyncio @@ -27,6 +26,7 @@ async def test_write_directory(language: str, topic: str): assert len(ret["directory"]) assert isinstance(ret["directory"][0], dict) + @pytest.mark.asyncio @pytest.mark.parametrize( ("language", "topic", "directory"), @@ -38,16 +38,3 @@ async def test_write_content(language: str, topic: str, directory: Dict): assert list(directory.keys())[0] in ret for value in list(directory.values())[0]: assert value in ret - -@pytest.mark.asyncio -@pytest.mark.parametrize( - ("title", "content"), - [("Python", "Write a tutorial about Python")] -) -async def test_save_docx(title: str, content: str): - ret = await SaveDocx().run(title=title, content=content) - assert isinstance(ret, str) - assert title in ret - async with aiofiles.open(ret, mode="r") as reader: - body = await reader.read() - assert body == content \ No newline at end of file diff --git a/tests/metagpt/utils/test_file.py b/tests/metagpt/utils/test_file.py new file mode 100644 index 000000000..a9f1a353d --- /dev/null +++ b/tests/metagpt/utils/test_file.py @@ -0,0 +1,27 @@ +#!/usr/bin/env python3 +# _*_ coding: utf-8 _*_ +""" +@Time : 2023/9/4 15:40:40 +@Author : Stitch-z +@File : test_file.py +""" +from pathlib import Path + +import aiofiles +import pytest + +from metagpt.utils.file import File + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("root_path", "filename", "content"), + [(Path("/code/MetaGPT/data/tutorial_docx/2023-09-07_17-05-20"), "test.md", "Hello World!")] +) +async def test_write_file(root_path: Path, filename: str, content: bytes): + full_file_name = await File.write(root_path=root_path, filename=filename, content=content.encode('utf-8')) + assert isinstance(full_file_name, Path) + assert root_path / filename == full_file_name + async with aiofiles.open(full_file_name, mode="r") as reader: + body = await reader.read() + assert body == content \ No newline at end of file