Multiagent.py

"""
STEW — Multi-Agent Coordinator
Orchestrates specialist sub-agents for complex multi-step tasks.
Each agent has a focused role; the Coordinator routes and sequences them.
"""
from openai import OpenAI
from config import Config


AGENT_ROLES = {
    "researcher": {
        "model": "owl-alpha",
        "system": (
            "You are a Research Agent. Your only job is to gather, verify, and summarise "
            "factual information. Search the web when needed. Be concise and cite sources."
        ),
    },
    "coder": {
        "model": "gpt-oss-120b",
        "system": (
            "You are a Senior Software Engineer Agent. Write clean, production-ready Python, "
            "FastAPI, React, or any requested stack. Include docstrings, error handling, and tests."
        ),
    },
    "planner": {
        "model": "owl-alpha",
        "system": (
            "You are a Task Planner Agent. Break complex goals into numbered, actionable steps. "
            "Output structured plans with clear milestones and dependencies."
        ),
    },
    "analyst": {
        "model": "gpt-oss-120b",
        "system": (
            "You are a Data Analyst Agent. Analyse data, identify patterns, create insights, "
            "and recommend decisions based on evidence."
        ),
    },
    "writer": {
        "model": "owl-alpha",
        "system": (
            "You are a Technical Writer Agent. Write documentation, reports, README files, "
            "and structured content. Be clear, professional, and thorough."
        ),
    },
    "reviewer": {
        "model": "gpt-oss-120b",
        "system": (
            "You are a Code Review Agent. Review code for bugs, security issues, performance "
            "problems, and style. Provide actionable improvement suggestions."
        ),
    },
}


class SubAgent:
    def __init__(self, role: str):
        if role not in AGENT_ROLES:
            raise ValueError(f"Unknown agent role: {role}. Choose from {list(AGENT_ROLES)}")
        self.role = role
        self.config = AGENT_ROLES[role]
        self.client = OpenAI(
            api_key=Config.OPENROUTER_API_KEY,
            base_url=Config.OPENROUTER_BASE_URL,
            default_headers={
                "HTTP-Referer": Config.OPENROUTER_SITE_URL,
                "X-Title": Config.OPENROUTER_APP_NAME,
            },
        )

    def run(self, task: str, context: str = "") -> str:
        messages = [{"role": "system", "content": self.config["system"]}]
        if context:
            messages.append({"role": "system", "content": f"Context from previous agents:\n{context}"})
        messages.append({"role": "user", "content": task})

        model_id = Config.MODELS.get(self.config["model"], Config.MODELS[Config.DEFAULT_MODEL])
        response = self.client.chat.completions.create(
            model=model_id,
            messages=messages,
            max_tokens=Config.MAX_RESPONSE_TOKENS,
        )
        return response.choices[0].message.content or ""


class MultiAgentCoordinator:
    """
    Routes tasks through multiple specialist agents and assembles results.
    """

    def __init__(self):
        self.agents = {role: SubAgent(role) for role in AGENT_ROLES}

    def run_pipeline(self, task: str, pipeline: list[str]) -> str:
        """
        Execute a sequential pipeline of agents.
        Each agent receives the original task + previous agents' outputs as context.

        Example pipeline: ["researcher", "planner", "coder", "reviewer"]
        """
        context = ""
        log = []

        for role in pipeline:
            if role not in self.agents:
                log.append(f"⚠️ Unknown agent: {role}, skipping.")
                continue

            log.append(f"\n{'='*50}\n🤖 {role.upper()} AGENT\n{'='*50}")
            result = self.agents[role].run(task, context)
            log.append(result)
            context += f"\n\n[{role.upper()} OUTPUT]:\n{result}"

        return "\n".join(log)

    def run_parallel(self, task: str, roles: list[str]) -> str:
        """Run multiple agents on the same task independently and combine results."""
        results = []
        for role in roles:
            if role in self.agents:
                output = self.agents[role].run(task)
                results.append(f"\n🤖 **{role.upper()}**:\n{output}")
        return "\n".join(results)

    def smart_route(self, task: str) -> str:
        """Auto-select the best pipeline based on task keywords."""
        task_lower = task.lower()

        if any(k in task_lower for k in ["code", "build", "create app", "write function", "fix bug"]):
            pipeline = ["planner", "coder", "reviewer"]
        elif any(k in task_lower for k in ["research", "find", "search", "who is", "what is"]):
            pipeline = ["researcher", "analyst"]
        elif any(k in task_lower for k in ["analyse", "analyze", "data", "report", "insight"]):
            pipeline = ["analyst", "writer"]
        elif any(k in task_lower for k in ["plan", "strategy", "steps", "how to", "roadmap"]):
            pipeline = ["planner", "writer"]
        else:
            pipeline = ["planner", "coder"]

        return self.run_pipeline(task, pipeline)

    def list_agents(self) -> str:
        lines = ["🤖 Available STEW Sub-Agents:\n"]
        for role, cfg in AGENT_ROLES.items():
            lines.append(f"• **{role}** — model: {cfg['model']}")
        return "\n".join(lines)