File size: 5,203 Bytes
553fbf7
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
"""
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