Software company simulation with corporate memory

from metagpt.roles import ProductManager, Architect, Engineer, QaEngineer
from metagpt.team import Team
import httpx
import os

class CompanyMemory:
    """Institutional memory for the AI software company."""

    def __init__(self):
        self.url = "https://mcp.stompy.ai/sse"
        self.headers = {"Authorization": f"Bearer {os.environ['STOMPY_TOKEN']}"}

    async def recall_decisions(self, domain: str) -> str:
        """Recall architectural decisions for a domain."""
        async with httpx.AsyncClient() as client:
            response = await client.post(self.url, headers=self.headers,
                json={"tool": "context_search", "query": f"architecture {domain}", "limit": 5})
            results = response.json().get("results", [])
            return "\n---\n".join([r["content"] for r in results])

    async def recall_patterns(self, pattern_type: str) -> str:
        """Recall established code patterns."""
        async with httpx.AsyncClient() as client:
            response = await client.post(self.url, headers=self.headers,
                json={"tool": "recall_context", "topic": f"patterns_{pattern_type}"})
            return response.json().get("content", "")

    async def save_decision(self, topic: str, content: str, priority: str = "important"):
        """Save an architectural decision for future reference."""
        async with httpx.AsyncClient() as client:
            await client.post(self.url, headers=self.headers,
                json={"tool": "lock_context", "topic": topic,
                      "content": content, "priority": priority})

company_memory = CompanyMemory()

class MemoryEnabledArchitect(Architect):
    """Architect role with access to institutional memory."""

    async def _think(self):
        # Recall relevant past decisions before designing
        past_decisions = await company_memory.recall_decisions(self.current_task.domain)
        existing_patterns = await company_memory.recall_patterns("architecture")

        self.context += f"""
INSTITUTIONAL KNOWLEDGE:
Past architectural decisions:
{past_decisions}

Established patterns:
{existing_patterns}

Design should be consistent with existing systems unless there's good reason to deviate.
"""
        return await super()._think()

    async def _act(self):
        result = await super()._act()

        # Save new architectural decisions
        if result.design_doc:
            await company_memory.save_decision(
                topic=f"adr_{self.current_task.name}",
                content=f"Design: {result.design_doc}\nRationale: {result.rationale}",
                priority="important"
            )

        return result

# Similarly for other roles
class MemoryEnabledPM(ProductManager):
    async def _think(self):
        past_prds = await company_memory.recall_decisions("requirements")
        self.context += f"\nPREVIOUS SUCCESSFUL PRDs:\n{past_prds}"
        return await super()._think()

async def run_project_with_memory(requirement: str, project_name: str):
    """Run a MetaGPT project with full institutional memory."""

    # Recall what the company knows about similar projects
    similar_projects = await company_memory.recall_decisions(f"project similar to {requirement}")

    team = Team(
        roles=[
            MemoryEnabledPM(),
            MemoryEnabledArchitect(),
            MemoryEnabledEngineer(),
            MemoryEnabledQA()
        ]
    )

    # Run the project
    result = await team.run(requirement)

    # Save project outcome for future learning
    await company_memory.save_decision(
        topic=f"project_{project_name}_outcome",
        content=f"""Project: {project_name}
Requirement: {requirement}
Architecture: {result.architecture_summary}
Key patterns used: {result.patterns}
Lessons learned: {result.retrospective}
Status: {result.status}""",
        priority="important"
    )

    return result

# Each project makes the company smarter
await run_project_with_memory("Build user authentication with OAuth", "auth_v1")
await run_project_with_memory("Add Stripe payment integration", "payments_v1")
# Second project now knows about the first!