From dd9a36f38883f9c830bb31cb2c0d4b0a5038d2c5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E8=8E=98=E6=9D=83=20=E9=A9=AC?= Date: Thu, 6 Jun 2024 12:10:27 +0800 Subject: [PATCH] feat: +additional output files --- examples/di/crawl_webpage.py | 3 +- metagpt/actions/design_api.py | 46 +++++++++++++----------- metagpt/actions/project_management.py | 12 +++---- metagpt/actions/write_prd.py | 28 ++++++++------- metagpt/rag/retrievers/bm25_retriever.py | 2 +- tests/metagpt/actions/test_design_api.py | 14 ++++---- tests/metagpt/actions/test_write_prd.py | 31 ++++++++-------- 7 files changed, 70 insertions(+), 66 deletions(-) diff --git a/examples/di/crawl_webpage.py b/examples/di/crawl_webpage.py index 10b230f2b..92e3c32b0 100644 --- a/examples/di/crawl_webpage.py +++ b/examples/di/crawl_webpage.py @@ -6,7 +6,8 @@ """ from metagpt.roles.di.data_interpreter import DataInterpreter -from metagpt.tools.libs.browser import Browser as _ + +__import__("metagpt.tools.libs.browser", fromlist=["Browser"]) # To skip pre-commit check PAPER_LIST_REQ = """" diff --git a/metagpt/actions/design_api.py b/metagpt/actions/design_api.py index 981dde53f..b4619f5fd 100644 --- a/metagpt/actions/design_api.py +++ b/metagpt/actions/design_api.py @@ -13,7 +13,7 @@ import json import uuid from pathlib import Path -from typing import List, Optional +from typing import List, Optional, Union from pydantic import BaseModel, Field @@ -70,7 +70,7 @@ class WriteDesign(Action): extra_info: str = "", output_pathname: str = "", **kwargs, - ) -> AIMessage: + ) -> Union[AIMessage, str]: """ Write a system design. @@ -90,7 +90,7 @@ class WriteDesign(Action): >>> extra_info = "Your extra information" >>> action = WriteDesign() >>> result = await action.run(user_requirement=user_requirement, extra_info=extra_info) - >>> print(result.content) + >>> print(result) System Design filename: "/path/to/design/filename" # Modify an exists system design. @@ -99,7 +99,7 @@ class WriteDesign(Action): >>> legacy_design_filename = "/path/to/exists/design/filename" >>> action = WriteDesign() >>> result = await action.run(user_requirement=user_requirement, extra_info=extra_info, legacy_design_filename=legacy_design_filename) - >>> print(result.content) + >>> print(result) System Design filename: "/path/to/design/filename" # Write a new system design with the given PRD(Product Requirement Document). @@ -108,7 +108,7 @@ class WriteDesign(Action): >>> prd_filename = "/path/to/prd/filename" >>> action = WriteDesign() >>> result = await action.run(user_requirement=user_requirement, extra_info=extra_info, prd_filename=prd_filename) - >>> print(result.content) + >>> print(result) System Design filename: "/path/to/design/filename" # Modify an exists system design with the given PRD(Product Requirement Document). @@ -118,7 +118,7 @@ class WriteDesign(Action): >>> legacy_design_filename = "/path/to/exists/design/filename" >>> action = WriteDesign() >>> result = await action.run(user_requirement=user_requirement, extra_info=extra_info, legacy_design_filename=legacy_design_filename, prd_filename=prd_filename) - >>> print(result.content) + >>> print(result) TSystem Design filename: "/path/to/design/filename" # Write a new system design and save to the path name. @@ -127,7 +127,7 @@ class WriteDesign(Action): >>> output_pathname = "/path/to/design/filename" >>> action = WriteDesign() >>> result = await action.run(user_requirement=user_requirement, extra_info=extra_info, output_pathname=output_pathname) - >>> print(result.content) + >>> print(result) System Design filename: "/path/to/design/filename" # Modify an exists system design and save to the path name. @@ -137,7 +137,7 @@ class WriteDesign(Action): >>> output_pathname = "/path/to/design/filename" >>> action = WriteDesign() >>> result = await action.run(user_requirement=user_requirement, extra_info=extra_info, legacy_design_filename=legacy_design_filename, output_pathname=output_pathname) - >>> print(result.content) + >>> print(result) System Design filename: "/path/to/design/filename" # Write a new system design with the given PRD(Product Requirement Document) and save to the path name. @@ -147,7 +147,7 @@ class WriteDesign(Action): >>> output_pathname = "/path/to/design/filename" >>> action = WriteDesign() >>> result = await action.run(user_requirement=user_requirement, extra_info=extra_info, prd_filename=prd_filename, output_pathname=output_pathname) - >>> print(result.content) + >>> print(result) System Design filename: "/path/to/design/filename" # Modify an exists system design with the given PRD(Product Requirement Document) and save to the path name. @@ -158,7 +158,7 @@ class WriteDesign(Action): >>> output_pathname = "/path/to/design/filename" >>> action = WriteDesign() >>> result = await action.run(user_requirement=user_requirement, extra_info=extra_info, legacy_design_filename=legacy_design_filename, prd_filename=prd_filename, output_pathname=output_pathname) - >>> print(result.content) + >>> print(result) System Design filename: "/path/to/design/filename" """ if not with_messages: @@ -241,21 +241,25 @@ class WriteDesign(Action): await reporter.async_report(self.repo.workdir / md.root_relative_path, "path") return doc - async def _save_data_api_design(self, design_doc): + async def _save_data_api_design(self, design_doc, output_filename: Path = None): m = json.loads(design_doc.content) data_api_design = m.get(DATA_STRUCTURES_AND_INTERFACES.key) or m.get(REFINED_DATA_STRUCTURES_AND_INTERFACES.key) if not data_api_design: return - pathname = self.repo.workdir / DATA_API_DESIGN_FILE_REPO / Path(design_doc.filename).with_suffix("") + pathname = output_filename or self.repo.workdir / DATA_API_DESIGN_FILE_REPO / Path( + design_doc.filename + ).with_suffix("") await self._save_mermaid_file(data_api_design, pathname) logger.info(f"Save class view to {str(pathname)}") - async def _save_seq_flow(self, design_doc): + async def _save_seq_flow(self, design_doc, output_filename: Path = None): m = json.loads(design_doc.content) seq_flow = m.get(PROGRAM_CALL_FLOW.key) or m.get(REFINED_PROGRAM_CALL_FLOW.key) if not seq_flow: return - pathname = self.repo.workdir / Path(SEQ_FLOW_FILE_REPO) / Path(design_doc.filename).with_suffix("") + pathname = output_filename or self.repo.workdir / Path(SEQ_FLOW_FILE_REPO) / Path( + design_doc.filename + ).with_suffix("") await self._save_mermaid_file(seq_flow, pathname) logger.info(f"Saving sequence flow to {str(pathname)}") @@ -273,7 +277,7 @@ class WriteDesign(Action): legacy_design_filename: str = "", extra_info: str = "", output_pathname: str = "", - ) -> AIMessage: + ) -> str: prd_content = "" if prd_filename: prd_content = await aread(filename=prd_filename) @@ -295,10 +299,10 @@ class WriteDesign(Action): output_path = DEFAULT_WORKSPACE_ROOT output_path.mkdir(parents=True, exist_ok=True) output_pathname = Path(output_path) / f"{uuid.uuid4().hex}.json" + output_pathname = Path(output_pathname) await awrite(filename=output_pathname, data=design.content) - kvs = {"changed_system_design_filenames": [output_pathname]} - - return AIMessage( - content=f'System Design filename: "{str(output_pathname)}"', - instruct_content=AIMessage.create_instruct_value(kvs=kvs), - ) + output_filename = output_pathname.parent / f"{output_pathname.stem}-class-diagram" + await self._save_data_api_design(design_doc=design, output_filename=output_filename) + output_filename = output_pathname.parent / f"{output_pathname.stem}-sequence-diagram" + await self._save_seq_flow(design_doc=design, output_filename=output_filename) + return f'System Design filename: "{str(output_pathname)}"' diff --git a/metagpt/actions/project_management.py b/metagpt/actions/project_management.py index b44bfb9f3..7196cf6c5 100644 --- a/metagpt/actions/project_management.py +++ b/metagpt/actions/project_management.py @@ -13,7 +13,7 @@ import json from pathlib import Path -from typing import List, Optional +from typing import List, Optional, Union from pydantic import BaseModel, Field @@ -45,7 +45,7 @@ class WriteTasks(Action): async def run( self, with_messages: List[Message] = None, *, user_requirement: str = "", design_filename: str = "", **kwargs - ) -> AIMessage: + ) -> Union[AIMessage, str]: """ Write a project schedule given a project system design file. @@ -62,7 +62,7 @@ class WriteTasks(Action): >>> design_filename = "/path/to/design/filename" >>> action = WriteTasks() >>> result = await action.run(design_filename=design_filename) - >>> print(result.content) + >>> print(result) The project schedule is balabala... # Write a new project schedule with the user requirement. @@ -70,7 +70,7 @@ class WriteTasks(Action): >>> user_requirement = "Your user requirements" >>> action = WriteTasks() >>> result = await action.run(design_filename=design_filename, user_requirement=user_requirement) - >>> print(result.content) + >>> print(result) The project schedule is balabala... """ if not with_messages: @@ -158,10 +158,10 @@ class WriteTasks(Action): packages.add(pkg) await self.repo.save(filename=PACKAGE_REQUIREMENTS_FILENAME, content="\n".join(packages)) - async def _execute_api(self, user_requirement: str = "", design_filename: str = ""): + async def _execute_api(self, user_requirement: str = "", design_filename: str = "") -> str: context = to_markdown_code_block(user_requirement) if not design_filename: content = await aread(filename=design_filename) context += to_markdown_code_block(content) node = await self._run_new_tasks(context) - return AIMessage(content=node.instruct_content.model_dump_json()) + return node.instruct_content.model_dump_json() diff --git a/metagpt/actions/write_prd.py b/metagpt/actions/write_prd.py index 5ed45cab4..0e4694423 100644 --- a/metagpt/actions/write_prd.py +++ b/metagpt/actions/write_prd.py @@ -17,7 +17,7 @@ from __future__ import annotations import json import uuid from pathlib import Path -from typing import List, Optional +from typing import List, Optional, Union from pydantic import BaseModel, Field @@ -87,7 +87,7 @@ class WritePRD(Action): legacy_prd_filename: str = "", extra_info: str = "", **kwargs, - ) -> AIMessage: + ) -> Union[AIMessage, str]: """ Write a Product Requirement Document. @@ -99,7 +99,7 @@ class WritePRD(Action): **kwargs: Additional keyword arguments. Returns: - AIMessage: The resulting message after generating the Product Requirement Document. + str: The resulting message after generating the Product Requirement Document. Example: # Write a new PRD(Product Requirement Document) @@ -107,7 +107,7 @@ class WritePRD(Action): >>> extra_info = "YOUR EXTRA INFO" >>> write_prd = WritePRD() >>> result = await write_prd.run(user_requirement=user_requirement, extra_info=extra_info) - >>> print(result.content) + >>> print(result) PRD filename: "/path/to/prd/directory/213434ad.json" # Modify a exists PRD(Product Requirement Document) @@ -116,7 +116,7 @@ class WritePRD(Action): >>> legacy_prd_filename = "/path/to/exists/prd_filename" >>> write_prd = WritePRD() >>> result = await write_prd.run(user_requirement=user_requirement, extra_info=extra_info, legacy_prd_filename=legacy_prd_filename) - >>> print(result.content) + >>> print(result) PRD filename: "/path/to/prd/directory/213434ad.json" # Write and save a new PRD(Product Requirement Document) to the path name. @@ -125,7 +125,7 @@ class WritePRD(Action): >>> output_pathname = "/path/to/prd/directory/213434ad.json" >>> write_prd = WritePRD() >>> result = await write_prd.run(user_requirement=user_requirement, extra_info=extra_info, output_pathname=output_pathname) - >>> print(result.content) + >>> print(result) PRD filename: "/path/to/prd/directory/213434ad.json" # Modify a exists PRD(Product Requirement Document) and save to the path name. @@ -135,7 +135,7 @@ class WritePRD(Action): >>> output_pathname = "/path/to/prd/directory/213434ad.json" >>> write_prd = WritePRD() >>> result = await write_prd.run(user_requirement=user_requirement, extra_info=extra_info, legacy_prd_filename=legacy_prd_filename, output_pathname=output_pathname) - >>> print(result.content) + >>> print(result) PRD filename: "/path/to/prd/directory/213434ad.json" """ @@ -201,7 +201,7 @@ class WritePRD(Action): cause_by=self, ) - async def _handle_bugfix(self, req: Document) -> Message: + async def _handle_bugfix(self, req: Document) -> AIMessage: # ... bugfix logic ... await self.repo.docs.save(filename=BUGFIX_FILENAME, content=req.content) await self.repo.docs.save(filename=REQUIREMENT_FILENAME, content="") @@ -283,12 +283,12 @@ class WritePRD(Action): await reporter.async_report(self.repo.workdir / md.root_relative_path, "path") return new_prd_doc - async def _save_competitive_analysis(self, prd_doc: Document): + async def _save_competitive_analysis(self, prd_doc: Document, output_filename: Path = None): m = json.loads(prd_doc.content) quadrant_chart = m.get(COMPETITIVE_QUADRANT_CHART.key) if not quadrant_chart: return - pathname = self.repo.workdir / COMPETITIVE_ANALYSIS_FILE_REPO / Path(prd_doc.filename).stem + pathname = output_filename or self.repo.workdir / COMPETITIVE_ANALYSIS_FILE_REPO / Path(prd_doc.filename).stem pathname.parent.mkdir(parents=True, exist_ok=True) await mermaid_to_file(self.config.mermaid.engine, quadrant_chart, pathname) image_path = pathname.parent / f"{pathname.name}.png" @@ -308,7 +308,7 @@ class WritePRD(Action): async def _execute_api( self, user_requirement: str, output_pathname: str, legacy_prd_filename: str, extra_info: str - ) -> AIMessage: + ) -> str: content = "#### User Requirements\n{user_requirement}\n#### Extra Info\n{extra_info}\n".format( user_requirement=to_markdown_code_block(val=user_requirement), extra_info=to_markdown_code_block(val=extra_info), @@ -326,6 +326,8 @@ class WritePRD(Action): output_path = DEFAULT_WORKSPACE_ROOT output_path.mkdir(parents=True, exist_ok=True) output_pathname = Path(output_path) / f"{uuid.uuid4().hex}.json" + output_pathname = Path(output_pathname) await awrite(filename=output_pathname, data=new_prd.content) - kvs = AIMessage.create_instruct_value({"changed_prd_filenames": [str(output_pathname)]}) - return AIMessage(content=f'PRD filename: "{str(output_pathname)}"', instruct_content=kvs) + competitive_analysis_filename = output_pathname.parent / f"{output_pathname.stem}-competitive-analysis" + await self._save_competitive_analysis(prd_doc=new_prd, output_filename=Path(competitive_analysis_filename)) + return f'PRD filename: "{str(output_pathname)}"' diff --git a/metagpt/rag/retrievers/bm25_retriever.py b/metagpt/rag/retrievers/bm25_retriever.py index 3b085cb73..dc75d87b0 100644 --- a/metagpt/rag/retrievers/bm25_retriever.py +++ b/metagpt/rag/retrievers/bm25_retriever.py @@ -46,4 +46,4 @@ class DynamicBM25Retriever(BM25Retriever): def persist(self, persist_dir: str, **kwargs) -> None: """Support persist.""" if self._index: - self._index.storage_context.persist(persist_dir) \ No newline at end of file + self._index.storage_context.persist(persist_dir) diff --git a/tests/metagpt/actions/test_design_api.py b/tests/metagpt/actions/test_design_api.py index 3398be5e6..314ff54e7 100644 --- a/tests/metagpt/actions/test_design_api.py +++ b/tests/metagpt/actions/test_design_api.py @@ -72,9 +72,9 @@ async def test_design_api(context, user_requirement, prd_filename, legacy_design result = await action.run( user_requirement=user_requirement, prd_filename=prd_filename, legacy_design_filename=legacy_design_filename ) - assert isinstance(result, AIMessage) - assert result.content - assert str(DEFAULT_WORKSPACE_ROOT) in result.content + assert isinstance(result, str) + assert result + assert str(DEFAULT_WORKSPACE_ROOT) in result @pytest.mark.parametrize( @@ -98,11 +98,9 @@ async def test_design_api_dir(context, user_requirement, prd_filename, legacy_de legacy_design_filename=legacy_design_filename, output_pathname=str(Path(context.config.project_path) / "1.txt"), ) - assert isinstance(result, AIMessage) - assert result.content - assert str(context.config.project_path) in result.content - assert result.instruct_content - assert result.instruct_content.changed_system_design_filenames + assert isinstance(result, str) + assert result + assert str(context.config.project_path) in result if __name__ == "__main__": diff --git a/tests/metagpt/actions/test_write_prd.py b/tests/metagpt/actions/test_write_prd.py index 6cf1da1dc..fcfa81931 100644 --- a/tests/metagpt/actions/test_write_prd.py +++ b/tests/metagpt/actions/test_write_prd.py @@ -16,7 +16,7 @@ from metagpt.const import DEFAULT_WORKSPACE_ROOT, REQUIREMENT_FILENAME from metagpt.logs import logger from metagpt.roles.product_manager import ProductManager from metagpt.roles.role import RoleReactMode -from metagpt.schema import AIMessage, Message +from metagpt.schema import Message from metagpt.utils.common import any_to_str from metagpt.utils.project_repo import ProjectRepo from tests.data.incremental_dev_project.mock import NEW_REQUIREMENT_SAMPLE @@ -79,35 +79,34 @@ async def test_fix_debug(new_filename, context, git_dir): async def test_write_prd_api(context): action = WritePRD() result = await action.run(user_requirement="write a snake game.") - assert isinstance(result, AIMessage) - assert result.content - assert str(DEFAULT_WORKSPACE_ROOT) in result.content + assert isinstance(result, str) + assert result + assert str(DEFAULT_WORKSPACE_ROOT) in result result = await action.run( user_requirement="write a snake game.", output_pathname=str(Path(context.config.project_path) / f"{uuid.uuid4().hex}.json"), ) - assert isinstance(result, AIMessage) - assert result.content - assert result.instruct_content - assert str(context.config.project_path) in result.content + assert isinstance(result, str) + assert result + assert str(context.config.project_path) in result - legacy_prd_filename = result.instruct_content.changed_prd_filenames[-1] + ix = result.find(":") + legacy_prd_filename = result[ix + 1 :].replace('"', "").strip() result = await action.run(user_requirement="Add moving enemy.", legacy_prd_filename=legacy_prd_filename) - assert isinstance(result, AIMessage) - assert result.content - assert str(DEFAULT_WORKSPACE_ROOT) in result.content + assert isinstance(result, str) + assert result + assert str(DEFAULT_WORKSPACE_ROOT) in result result = await action.run( user_requirement="Add moving enemy.", output_pathname=str(Path(context.config.project_path) / f"{uuid.uuid4().hex}.json"), legacy_prd_filename=legacy_prd_filename, ) - assert isinstance(result, AIMessage) - assert result.content - assert result.instruct_content - assert str(context.config.project_path) in result.content + assert isinstance(result, str) + assert result + assert str(context.config.project_path) in result if __name__ == "__main__":