src/agents/orchestrator.py

"""
MINDI 1.5 Vision-Coder — Agent Orchestrator

Coordinates multiple AI agents (Code Gen, Vision Critic, Search, Sandbox)
to produce, evaluate, and refine generated code.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from enum import Enum
from pathlib import Path
from typing import Any, Optional


class AgentRole(str, Enum):
    """Roles for MINDI's agent system."""
    CODE_GENERATOR = "code_generator"
    UI_CRITIC = "ui_critic"
    SEARCH_AGENT = "search_agent"
    SANDBOX_RUNNER = "sandbox_runner"
    ERROR_FIXER = "error_fixer"


@dataclass
class AgentMessage:
    """A message passed between agents in the orchestration pipeline."""
    role: AgentRole
    content: str
    metadata: dict[str, Any] = field(default_factory=dict)
    artifacts: list[Path] = field(default_factory=list)


@dataclass
class GenerationResult:
    """Final output from the agent pipeline."""
    code: str
    language: str
    file_path: str
    critique: Optional[str] = None
    search_context: Optional[str] = None
    sandbox_output: Optional[str] = None
    iterations: int = 1
    success: bool = True
    errors: list[str] = field(default_factory=list)


class AgentOrchestrator:
    """
    Orchestrates the MINDI agent pipeline:

    1. User prompt arrives
    2. Search Agent gathers relevant docs/packages
    3. Code Generator produces Next.js + Tailwind + TS code
    4. Sandbox Runner tests the code in isolation
    5. Vision Critic screenshots the output and evaluates UI/UX
    6. Error Fixer resolves any issues
    7. Loop until quality threshold or max iterations
    """

    def __init__(
        self,
        max_iterations: int = 3,
        quality_threshold: float = 0.85,
        log_dir: Optional[Path] = None,
    ) -> None:
        self.max_iterations = max_iterations
        self.quality_threshold = quality_threshold
        self.log_dir = log_dir or Path("./logs/agents")
        self.log_dir.mkdir(parents=True, exist_ok=True)
        self.history: list[AgentMessage] = []

    async def run_pipeline(
        self,
        user_prompt: str,
        context: Optional[dict[str, Any]] = None,
    ) -> GenerationResult:
        """
        Execute the full agent pipeline for a user request.

        This is the main entry point — called by the FastAPI backend.
        Each step will be implemented as we build each agent module.
        """
        self.history.clear()
        context = context or {}

        # Step 1: Search for relevant documentation
        search_result = await self._run_search(user_prompt)

        # Step 2: Generate code
        code_result = await self._generate_code(user_prompt, search_result)

        # Step 3: Test in sandbox
        sandbox_result = await self._run_sandbox(code_result)

        # Step 4: Vision critique (if sandbox produced a screenshot)
        critique_result = await self._run_critique(code_result, sandbox_result)

        # Step 5: Fix errors if any
        final_code = code_result
        iterations = 1

        while iterations < self.max_iterations:
            if sandbox_result.get("success") and critique_result.get("score", 0) >= self.quality_threshold:
                break
            final_code = await self._fix_errors(
                final_code, sandbox_result, critique_result
            )
            sandbox_result = await self._run_sandbox(final_code)
            critique_result = await self._run_critique(final_code, sandbox_result)
            iterations += 1

        return GenerationResult(
            code=final_code,
            language="typescript",
            file_path="page.tsx",
            critique=critique_result.get("feedback"),
            search_context=search_result.get("context"),
            sandbox_output=sandbox_result.get("output"),
            iterations=iterations,
            success=sandbox_result.get("success", False),
        )

    async def _run_search(self, prompt: str) -> dict[str, Any]:
        """Search for relevant docs and packages. Implemented in src/search/."""
        # Placeholder — will be wired to SearchAgent
        return {"context": "", "sources": []}

    async def _generate_code(self, prompt: str, search_ctx: dict[str, Any]) -> str:
        """Generate code using the fine-tuned model. Implemented in src/inference/."""
        # Placeholder — will be wired to inference pipeline
        return ""

    async def _run_sandbox(self, code: str) -> dict[str, Any]:
        """Run code in sandbox. Implemented in src/sandbox/."""
        # Placeholder — will be wired to SandboxRunner
        return {"success": False, "output": "", "screenshot": None}

    async def _run_critique(self, code: str, sandbox: dict[str, Any]) -> dict[str, Any]:
        """Critique UI via vision. Implemented in src/agents/ui_critic.py."""
        # Placeholder — will be wired to VisionCritic
        return {"score": 0.0, "feedback": ""}

    async def _fix_errors(
        self, code: str, sandbox: dict[str, Any], critique: dict[str, Any]
    ) -> str:
        """Fix errors in code. Implemented in src/agents/error_fixer.py."""
        # Placeholder — will be wired to ErrorFixer
        return code