agent/roles/software_company.py

import json
import os
from pathlib import Path

import aiofiles
from aiobotocore.session import get_session
from mdutils.mdutils import MdUtils
from metagpt.actions import Action, WriteCodeReview
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 (
    CODE_SUMMARIES_FILE_REPO,
    COMPETITIVE_ANALYSIS_FILE_REPO,
    DATA_API_DESIGN_FILE_REPO,
    DOCS_FILE_REPO,
    PRDS_FILE_REPO,
    REQUIREMENT_FILENAME,
    SEQ_FLOW_FILE_REPO,
    SYSTEM_DESIGN_FILE_REPO,
    TASK_FILE_REPO,
)
from metagpt.environment import Environment
from metagpt.roles import Architect, Engineer, ProductManager, ProjectManager, Role
from metagpt.schema import Message, MessageQueue
from metagpt.team import Team
from metagpt.utils.common import any_to_str, is_subscribed
from metagpt.utils.file_repository import FileRepository
from pydantic import Field
from zipstream import AioZipStream


class RoleRun(Action):
    role: Role
    desc: str  # f"{role.profile} {role.todo}"


class SoftwareCompanyEnv(Environment):
    buf: MessageQueue = Field(default_factory=MessageQueue)
    peek_buf: MessageQueue = Field(default_factory=MessageQueue)

    def publish_message(self, message: Message, peekable: bool = True) -> bool:
        self.buf.push(msg=message)
        if peekable:
            self.peek_buf.push(msg=message)
        return True

    def peek_messages(self):
        return self.peek_buf.pop_all()

    def get_message(self):
        return self.buf.pop()


class SoftwareCompany(Role):
    """封装软件公司成角色，以快速接入agent store。"""

    company: Team = Field(default_factory=Team)

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.company.env = SoftwareCompanyEnv()
        self.company.hire([ProductManager(), Architect(), ProjectManager(), Engineer(n_borg=5)])

    async def _think(self) -> bool:
        """软件公司运行需要4轮

        BOSS            -> ProductManager -> Architect   -> ProjectManager -> Engineer
        BossRequirement -> WritePRD       -> WriteDesign -> WriteTasks     -> WriteCode
        """
        while msg := self.company.env.get_message():
            for role in self.company.env.roles.values():
                if not is_subscribed(message=msg, tags=role.subscription) or not role.is_watch(caused_by=msg.cause_by):
                    continue
                role.put_message(message=msg)
                self.rc.todo = RoleRun(role=role, desc=f"{role.profile} {role.todo}")
                return True
        self.rc.todo = None
        return False

    async def _act(self) -> Message:
        await self.company.run(1, auto_archive=False)
        msgs = self.company.env.peek_messages()
        mappings = {
            any_to_str(PrepareDocuments): self.format_requirements,
            any_to_str(WritePRD): self.format_prd,
            any_to_str(WriteDesign): self.format_system_design,
            any_to_str(WriteTasks): self.format_task,
            any_to_str(WriteCode): self.format_code,
            any_to_str(WriteCodeReview): self.format_code,
            any_to_str(SummarizeCode): self.format_summarize_code,
        }
        for msg in msgs:
            func = mappings.get(msg.cause_by)
            if not func:
                continue
            output = await func(msg)
            return output

    async def format_requirements(self, msg: Message):
        mdfile = MdUtils("")
        doc = await FileRepository.get_file(filename=REQUIREMENT_FILENAME, relative_path=DOCS_FILE_REPO)
        mdfile.new_header(2, "New Requirements", add_table_of_contents="n")
        mdfile.new_paragraph(text=doc.content)
        return Message(mdfile.get_md_text(), cause_by=msg.cause_by, role=msg.role)

    async def format_summarize_code(self, msg: Message):
        mdfile = MdUtils("")
        docs = await FileRepository.get_all_files(relative_path=CODE_SUMMARIES_FILE_REPO)
        if not docs:
            mdfile.new_header(2, "Code Summaries", add_table_of_contents="n")
            mdfile.new_paragraph(text="Pass.")
            return Message(mdfile.get_md_text(), cause_by=msg.cause_by, role=msg.role)

        ix = 0
        for doc in docs:
            if ix > 0:
                mdfile.new_line("---")
            ix += 1
            mdfile.new_header(1, "RID:" + str(Path(doc.filename).with_suffix("")), add_table_of_contents="n")
            data = json.loads(doc.content)
            mdfile.new_header(2, "File names", add_table_of_contents="n")
            mdfile.new_list(data.get("codes_filenames", []))

            mdfile.new_header(2, "Summary", add_table_of_contents="n")
            mdfile.insert_code(data.get("reason", ""), "json")
        return Message(mdfile.get_md_text(), cause_by=msg.cause_by, role=msg.role)

    async def format_prd(self, prd: Message):
        mdfile = MdUtils("")
        docs = await FileRepository.get_all_files(relative_path=PRDS_FILE_REPO)
        ix = 0
        for doc in docs:
            filename = Path(doc.filename)
            if ix > 0:
                mdfile.new_line("---")
            ix += 1
            mdfile.new_header(1, "RID:" + str(filename.with_suffix("")), add_table_of_contents="n")

            data = json.loads(doc.content)
            for title, content in data.items():
                if title == "Competitive Analysis":
                    mdfile.new_header(2, title, add_table_of_contents="n")
                    if all(i.count(":") == 1 for i in content):
                        mdfile.new_table(
                            2,
                            len(content) + 1,
                            ["Competitor", "Description", *(i for j in content for i in j.split(":"))],
                        )
                    else:
                        mdfile.new_list(content, marked_with="1")
                elif title == "Competitive Quadrant Chart":
                    mdfile.new_header(2, title, add_table_of_contents="n")
                    competitive_analysis_path = (
                        CONFIG.git_repo.workdir / COMPETITIVE_ANALYSIS_FILE_REPO / filename.with_suffix(".png")
                    )
                    if not await self._add_s3_url(md_file=mdfile, title=title, pathname=competitive_analysis_path):
                        mdfile.insert_code(content, "mermaid")
                elif title == "Requirement Pool":
                    mdfile.new_header(2, title, add_table_of_contents="n")
                    mdfile.new_table(
                        2, len(content) + 1, ["Task Description", "Priority", *(i for j in content for i in j)]
                    )
                elif isinstance(content, list):
                    mdfile.new_header(2, title, add_table_of_contents="n")
                    mdfile.new_list(content, marked_with="1")
                else:
                    mdfile.new_header(2, title, add_table_of_contents="n")
                    mdfile.new_paragraph(content)
        return Message(mdfile.get_md_text(), cause_by=prd.cause_by, role=prd.role)

    async def format_system_design(self, design: Message):
        mdfile = MdUtils("")
        docs = await FileRepository.get_all_files(relative_path=SYSTEM_DESIGN_FILE_REPO)
        ix = 0
        for doc in docs:
            filename = Path(doc.filename)
            if ix > 0:
                mdfile.new_line("---")
            ix += 1
            mdfile.new_header(1, "RID:" + str(filename.with_suffix("")), add_table_of_contents="n")

            data = json.loads(doc.content)
            for title, content in data.items():
                if title == "Project name":
                    mdfile.new_header(2, title, add_table_of_contents="n")
                    mdfile.insert_code(content, "python")
                elif title == "Data structures and interfaces":
                    mdfile.new_header(2, title, add_table_of_contents="n")
                    data_api_design_path = (
                        CONFIG.git_repo.workdir / DATA_API_DESIGN_FILE_REPO / filename.with_suffix(".png")
                    )
                    if not await self._add_s3_url(md_file=mdfile, title=title, pathname=data_api_design_path):
                        mdfile.insert_code(content, "mermaid")
                elif title == "Program call flow":
                    mdfile.new_header(2, title, add_table_of_contents="n")
                    seq_flow_path = CONFIG.git_repo.workdir / SEQ_FLOW_FILE_REPO / filename.with_suffix(".png")
                    if not await self._add_s3_url(md_file=mdfile, title=title, pathname=seq_flow_path):
                        mdfile.insert_code(content, "mermaid")
                elif isinstance(content, list):
                    mdfile.new_header(2, title, add_table_of_contents="n")
                    mdfile.new_list(content, marked_with="1")
                else:
                    mdfile.new_header(2, title, add_table_of_contents="n")
                    mdfile.new_paragraph(content)
        return Message(mdfile.get_md_text(), cause_by=design.cause_by, role=design.role)

    async def format_task(self, task: Message):
        mdfile = MdUtils("")
        docs = await FileRepository.get_all_files(relative_path=TASK_FILE_REPO)
        ix = 0
        for doc in docs:
            filename = Path(doc.filename)
            if ix > 0:
                mdfile.new_line("---")
            ix += 1
            mdfile.new_header(1, "RID:" + str(filename.with_suffix("")), add_table_of_contents="n")

            data = json.loads(doc.content)
            for title, content in data.items():
                if title in {
                    "Required Python third-party packages",
                    "Required Other language third-party packages",
                    "Task list",
                }:
                    mdfile.new_header(2, title, add_table_of_contents="n")
                    mdfile.new_list(content)
                elif title == "Full API spec":
                    mdfile.new_header(2, title, add_table_of_contents="n")
                    mdfile.insert_code(content, "text")
                elif title == "Logic Analysis":
                    mdfile.new_header(2, title, add_table_of_contents="n")
                    mdfile.new_table(
                        2, len(content) + 1, ["Filename", "Class/Function Name", *(i for j in content for i in j)]
                    )
                elif title == "Shared Knowledge":
                    mdfile.new_header(2, title=title, add_table_of_contents="n")
                    mdfile.insert_code(content, "text")
                elif isinstance(content, list):
                    mdfile.new_header(2, title, add_table_of_contents="n")
                    mdfile.new_list(content, marked_with="1")
                else:
                    mdfile.new_header(2, title=title, add_table_of_contents="n")
                    mdfile.new_paragraph(text=content)
        return Message(mdfile.get_md_text(), cause_by=task.cause_by, role=task.role)

    async def format_code(self, code: Message):
        mdfile = MdUtils("")
        docs = await FileRepository.get_all_files(relative_path=CONFIG.src_workspace)
        code_block_types = {
            ".py": "python",
            ".yaml": "yaml",
            ".yml": "yaml",
            ".json": "json",
            ".js": "javascript",
            ".sql": "sql",
        }
        for doc in docs:
            mdfile.new_header(2, doc.filename, add_table_of_contents="n")
            suffix = doc.filename.rsplit(".", maxsplit=1)[-1]
            mdfile.insert_code(doc.content, code_block_types.get(suffix, "text"))
        url = await self.upload()
        mdfile.new_header(2, "Project Packaging Complete", add_table_of_contents="n")
        mdfile.new_paragraph(
            "We are thrilled to inform you that our project has been successfully packaged "
            "and is ready for download and use. You can download the packaged project through"
            f" the following link:\n[Project Download Link]({url})"
        )
        return Message(mdfile.get_md_text(), cause_by=code.cause_by, role=code.role)

    async def upload_file_to_s3(self, filepath: str, key: str):
        async with aiofiles.open(filepath, "rb") as f:
            content = await f.read()
            return await self.upload_to_s3(content, key)

    async def upload_to_s3(self, content: bytes, key: str):
        session = get_session()
        async with session.create_client(
            "s3",
            aws_secret_access_key=os.getenv("S3_SECRET_KEY"),
            aws_access_key_id=os.getenv("S3_ACCESS_KEY"),
            endpoint_url=os.getenv("S3_ENDPOINT_URL"),
            use_ssl=os.getenv("S3_SECURE"),
        ) as client:
            # upload object to amazon s3
            bucket = os.getenv("S3_BUCKET")
            await client.put_object(Bucket=bucket, Key=key, Body=content)
            return f"{os.getenv('S3_ENDPOINT_URL')}/{bucket}/{key}"

    async def upload(self):
        files = []
        all_filenames = CONFIG.git_repo.get_files(relative_path=".")
        for filename in all_filenames:
            full_path = CONFIG.git_repo.workdir / filename
            files.append({"file": str(full_path), "name": str(filename)})
        # aiozipstream
        chunks = []
        async for chunk in AioZipStream(files, chunksize=32768).stream():
            chunks.append(chunk)
        key = f"{CONFIG.workspace_path.name}/metagpt-{CONFIG.git_repo.workdir.name}.zip"
        return await self.upload_to_s3(b"".join(chunks), key)

    async def _add_s3_url(self, md_file: MdUtils, title: str, pathname: Path) -> bool:
        if pathname.exists():
            key = pathname.relative_to(CONFIG.workspace_path.parent)
            url = await self.upload_file_to_s3(filepath=str(pathname), key=str(key))
            md_file.new_line(md_file.new_inline_image(title, url))
            return True
        return False