software_company.py

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