MetaGPT-SoftwareCompany / software_company.py
shenchucheng
create meatgpt webui
faf9dc6
raw
history blame
14.6 kB
import asyncio
import json
from pydantic import BaseModel, Field
import datetime
import os
from pathlib import Path
from typing import Any, Coroutine, Optional
import aiofiles
from aiobotocore.session import get_session
from mdutils.mdutils import MdUtils
from zipstream import AioZipStream
from metagpt.actions import Action
from metagpt.actions.action_output import ActionOutput
from metagpt.actions.design_api import WriteDesign
from metagpt.actions.prepare_documents import PrepareDocuments
from metagpt.actions.project_management import WriteTasks
from metagpt.actions.summarize_code import SummarizeCode
from metagpt.actions.write_code import WriteCode
from metagpt.actions.write_prd import WritePRD
from metagpt.config import CONFIG
from metagpt.const import COMPETITIVE_ANALYSIS_FILE_REPO, DATA_API_DESIGN_FILE_REPO, SEQ_FLOW_FILE_REPO, SERDESER_PATH
from metagpt.roles import Architect, Engineer, ProductManager, ProjectManager, Role
from metagpt.schema import Message
from metagpt.team import Team
from metagpt.utils.common import any_to_str, read_json_file, write_json_file
from metagpt.utils.git_repository import GitRepository
class RoleRun(Action):
role: Role
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
action = self.role._rc.todo
self.desc = f"{self.role.profile} {action.desc or str(action)}"
class PackProject(Action):
role: Role
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.desc = "Pack the project with prd, design, code and more."
async def run(self, key: str):
url = await self.upload(key)
mdfile = MdUtils(None)
mdfile.new_line(mdfile.new_inline_link(url, url.rsplit("/", 1)[-1]))
return ActionOutput(mdfile.get_md_text(), BaseModel())
async def upload(self, key: str):
files = []
workspace = CONFIG.git_repo.workdir
workspace = str(workspace)
for r, _, fs in os.walk(workspace):
_r = r[len(workspace):].lstrip("/")
for f in fs:
files.append({"file": os.path.join(r, f), "name": os.path.join(_r, f)})
# aiozipstream
chunks = []
async for chunk in AioZipStream(files, chunksize=32768).stream():
chunks.append(chunk)
return await upload_to_s3(b"".join(chunks), key)
class SoftwareCompany(Role):
"""封装软件公司成角色,以快速接入agent store。"""
finish: bool = False
company: Team = Field(default_factory=Team)
active_role: Optional[Role] = None
git_repo: Optional[GitRepository] = None
max_auto_summarize_code: int = 0
def __init__(self, use_code_review=False, *args, **kwargs):
super().__init__(*args, **kwargs)
engineer = Engineer(n_borg=5, use_code_review=use_code_review)
self.company.hire([ProductManager(), Architect(), ProjectManager(), engineer])
self._init_actions([PackProject(role=engineer)])
def recv(self, message: Message) -> None:
self.company.run_project(message.content)
async def _think(self) -> Coroutine[Any, Any, bool]:
"""软件公司运行需要4轮
BOSS -> ProductManager -> Architect -> ProjectManager -> Engineer
BossRequirement -> WritePRD -> WriteDesign -> WriteTasks -> WriteCode ->
"""
if self.finish:
self._rc.todo = None
return False
if self.git_repo is not None:
CONFIG.git_repo = self.git_repo
environment = self.company.env
for role in environment.roles.values():
if await role._observe():
await role._think()
if isinstance(role._rc.todo, PrepareDocuments):
self.active_role = role
await self.act()
self.git_repo = CONFIG.git_repo
return await self._think()
if isinstance(role._rc.todo, SummarizeCode):
return await self._think()
self._rc.todo = RoleRun(role=role)
self.active_role = role
return True
self._set_state(0)
return True
async def _act(self) -> Message:
if self.git_repo is not None:
CONFIG.git_repo = self.git_repo
CONFIG.src_workspace = CONFIG.git_repo.workdir / CONFIG.git_repo.workdir.name
CONFIG.max_auto_summarize_code = self.max_auto_summarize_code
if isinstance(self._rc.todo, PackProject):
workdir = CONFIG.git_repo.workdir
name = workdir.name
uid = workdir.parent.parent.name
now = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
key = f"{uid}/metagpt-{name}-{now}.zip"
output = await self._rc.todo.run(key)
self.finish = True
return Message(output.content, role=self.profile, cause_by=type(self._rc.todo))
output = await self.active_role._act()
self.active_role._set_state(state=-1)
self.active_role.publish_message(output)
cause_by = output.cause_by
if cause_by == any_to_str(WritePRD):
output = await self.format_prd(output)
elif cause_by == any_to_str(WriteDesign):
output = await self.format_system_design(output)
elif cause_by == any_to_str(WriteTasks):
output = await self.format_tasks(output)
elif cause_by == any_to_str(WriteCode):
output = await self.format_code(output)
elif cause_by == any_to_str(SummarizeCode):
output = await self.format_code_summary(output)
return output
async def format_prd(self, msg: Message):
docs = [(k, v) for k, v in msg.instruct_content.docs.items()]
prd_doc = docs[0][1]
data = json.loads(prd_doc.content)
mdfile = MdUtils(None)
title = "Original Requirements"
mdfile.new_header(2, title, add_table_of_contents=False)
mdfile.new_paragraph(data[title])
title = "Product Goals"
mdfile.new_header(2, title, add_table_of_contents=False)
mdfile.new_list(data[title], marked_with="1")
title = "User Stories"
mdfile.new_header(2, title, add_table_of_contents=False)
mdfile.new_list(data[title], marked_with="1")
title = "Competitive Analysis"
mdfile.new_header(2, title, add_table_of_contents=False)
if all(i.count(":") == 1 for i in data[title]):
mdfile.new_table(
2, len(data[title]) + 1, ["Competitor", "Description", *(i for j in data[title] for i in j.split(":"))]
)
else:
mdfile.new_list(data[title], marked_with="1")
title = "Competitive Quadrant Chart"
mdfile.new_header(2, title, add_table_of_contents=False)
competitive_analysis_path = CONFIG.git_repo.workdir / Path(COMPETITIVE_ANALYSIS_FILE_REPO) / Path(prd_doc.filename).with_suffix(".png")
if competitive_analysis_path.exists():
key = str(competitive_analysis_path.relative_to(CONFIG.git_repo.workdir.parent.parent))
url = await upload_file_to_s3(competitive_analysis_path, key)
mdfile.new_line(mdfile.new_inline_image(title, url))
else:
mdfile.insert_code(data[title], "mermaid")
title = "Requirement Analysis"
mdfile.new_header(2, title, add_table_of_contents=False)
mdfile.new_paragraph(data[title])
title = "Requirement Pool"
mdfile.new_header(2, title, add_table_of_contents=False)
mdfile.new_table(
2, len(data[title]) + 1, ["Task Description", "Priority", *(i for j in data[title] for i in j)]
)
title = "UI Design draft"
mdfile.new_header(2, title, add_table_of_contents=False)
mdfile.new_paragraph(data[title])
title = "Anything UNCLEAR"
mdfile.new_header(2, title, add_table_of_contents=False)
mdfile.new_paragraph(data[title])
content = mdfile.get_md_text()
return Message(content, cause_by=msg.cause_by, role=msg.role)
async def format_system_design(self, msg: Message):
system_designs = [(k, v) for k, v in msg.instruct_content.docs.items()]
system_design_doc = system_designs[0][1]
data = json.loads(system_design_doc.content)
mdfile = MdUtils(None)
title = "Implementation approach"
mdfile.new_header(2, title, add_table_of_contents=False)
mdfile.new_paragraph(data[title])
title = "File list"
mdfile.new_header(2, title, add_table_of_contents=False)
mdfile.new_list(data[title], marked_with="1")
title = "Data structures and interfaces"
mdfile.new_header(2, title, add_table_of_contents=False)
data_api_design_path = CONFIG.git_repo.workdir / Path(DATA_API_DESIGN_FILE_REPO) / Path(system_design_doc.filename).with_suffix(".png")
if data_api_design_path.exists():
key = str(data_api_design_path.relative_to(CONFIG.git_repo.workdir.parent.parent))
url = await upload_file_to_s3(data_api_design_path, key)
mdfile.new_line(mdfile.new_inline_image(title, url))
else:
mdfile.insert_code(data[title], "mermaid")
title = "Program call flow"
mdfile.new_header(2, title, add_table_of_contents=False)
seq_flow_path = CONFIG.git_repo.workdir / SEQ_FLOW_FILE_REPO / Path(system_design_doc.filename).with_suffix(".png")
if seq_flow_path.exists():
key = str(seq_flow_path.relative_to(CONFIG.git_repo.workdir.parent.parent))
url = await upload_file_to_s3(seq_flow_path, key)
mdfile.new_line(mdfile.new_inline_image(title, url))
else:
mdfile.insert_code(data[title], "mermaid")
title = "Anything UNCLEAR"
mdfile.new_header(2, title, add_table_of_contents=False)
mdfile.new_paragraph(data[title])
content = mdfile.get_md_text()
return Message(content, cause_by=msg.cause_by, role=msg.role)
async def format_tasks(self, msg: Message):
tasks = [(k, v) for k, v in msg.instruct_content.docs.items()]
task_doc = tasks[0][1]
data = json.loads(task_doc.content)
mdfile = MdUtils(None)
title = "Required Python packages"
mdfile.new_header(2, title, add_table_of_contents=False)
mdfile.insert_code("\n".join(data[title]), "txt")
title = "Required Other language third-party packages"
mdfile.new_header(2, title, add_table_of_contents=False)
mdfile.insert_code("\n".join(data[title]), "txt")
title = "Logic Analysis"
mdfile.new_header(2, title, add_table_of_contents=False)
mdfile.new_table(
2, len(data[title]) + 1, ["Filename", "Class/Function Name", *(i for j in data[title] for i in j)]
)
title = "Task list"
mdfile.new_header(2, title, add_table_of_contents=False)
mdfile.new_list(data[title])
title = "Full API spec"
mdfile.new_header(2, title, add_table_of_contents=False)
if data[title]:
mdfile.insert_code(data[title], "json")
title = "Shared Knowledge"
mdfile.new_header(2, title, add_table_of_contents=False)
mdfile.insert_code(data[title], "python")
title = "Anything UNCLEAR"
mdfile.new_header(2, title, add_table_of_contents=False)
mdfile.insert_code(data[title], "python")
content = mdfile.get_md_text()
return Message(content, cause_by=msg.cause_by, role=msg.role)
async def format_code(self, msg: Message):
data = msg.content.splitlines()
workdir = CONFIG.git_repo.workdir
code_root = workdir / workdir.name
mdfile = MdUtils(None)
for filename in data:
mdfile.new_header(2, filename, add_table_of_contents=False)
async with aiofiles.open(code_root / filename) as f:
content = await f.read()
suffix = filename.rsplit(".", maxsplit=1)[-1]
mdfile.insert_code(content, "python" if suffix == "py" else suffix)
return Message(mdfile.get_md_text(), cause_by=msg.cause_by, role=msg.role)
async def format_code_summary(self, msg: Message):
# TODO
return msg
async def think(self):
await self._think()
return self._rc.todo
async def act(self):
return await self._act()
def serialize(self, stg_path: Path = None):
stg_path = SERDESER_PATH.joinpath("software_company") if stg_path is None else stg_path
team_info_path = stg_path.joinpath("software_company_info.json")
write_json_file(team_info_path, self.dict(exclude={"company": True}))
self.company.serialize(stg_path.joinpath("company")) # save company alone
@classmethod
def deserialize(cls, stg_path: Path) -> "Team":
"""stg_path = ./storage/team"""
# recover team_info
software_company_info_path = stg_path.joinpath("software_company_info.json")
if not software_company_info_path.exists():
raise FileNotFoundError(
"recover storage meta file `team_info.json` not exist, "
"not to recover and please start a new project."
)
software_company_info: dict = read_json_file(software_company_info_path)
# recover environment
company = Team.deserialize(stg_path=stg_path.joinpath("company"))
software_company_info.update({"company": company})
return cls(**software_company_info)
async def upload_file_to_s3(filepath: str, key: str):
async with aiofiles.open(filepath, "rb") as f:
content = await f.read()
return await upload_to_s3(content, key)
async def upload_to_s3(content: bytes, key: str):
session = get_session()
async with session.create_client(
"s3",
aws_secret_access_key=CONFIG.get("S3_SECRET_KEY"),
aws_access_key_id=CONFIG.get("S3_ACCESS_KEY"),
endpoint_url=CONFIG.get("S3_ENDPOINT_URL"),
use_ssl=CONFIG.get("S3_SECURE"),
) as client:
# upload object to amazon s3
bucket = CONFIG.get("S3_BUCKET")
await client.put_object(Bucket=bucket, Key=key, Body=content)
return f"{CONFIG.get('S3_ENDPOINT_URL')}/{bucket}/{key}"
async def main(idea, **kwargs):
sc = SoftwareCompany(**kwargs)
sc.recv(Message(idea))
while await sc.think():
print(await sc.act())
if __name__ == "__main__":
asyncio.run(main())