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 get_download_url(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 get_download_url(content, key) async def get_download_url(content: bytes, key: str) -> str: if CONFIG.get("STORAGE_TYPE") == "S3": 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}" else: storage = CONFIG.get("LOCAL_ROOT", "storage") base_url = CONFIG.get("LOCAL_BASE_URL", "storage") filepath = Path(storage) / key filepath.parent.mkdir(exist_ok=True, parents=True) async with aiofiles.open(filepath, "wb") as f: await f.write(content) return f"{base_url}/{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())