feat: +additional output files

This commit is contained in:
莘权 马 2024-06-06 12:10:27 +08:00
parent f37b12a8ec
commit dd9a36f388
7 changed files with 70 additions and 66 deletions

View file

@ -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 = """"

View file

@ -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)}"'

View file

@ -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()

View file

@ -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)}"'

View file

@ -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)
self._index.storage_context.persist(persist_dir)

View file

@ -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__":

View file

@ -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__":