Update app.py

app.py

@@ -1,231 +1,219 @@
 import gradio as gr
-import requests
 import json
 import re
 from typing import Dict, List, Tuple
 import os
 
-# Let's use a model that's good at this task and very cheap/fast
-SELECTED_MODEL = "minimax/minimax-m2:free"
+from crewai import Agent, Task, Crew, Process
+from langchain_openai import ChatOpenAI
+
+# --- Configuration ---
+# Set a default model. Using a fast and capable one is best.
+# Claude 3 Haiku is excellent for this and very cheap.
+DEFAULT_MODEL = "minimax/minimax-m2:free"
+
+# --- Agent Definitions ---
+
+# 1. Color Specialist Agent
+color_specialist = Agent(
+    role='Expert Color Palette Designer',
+    goal="""Select a visually appealing and contextually appropriate color palette 
+    based on the user's website description. The palette must include primary, 
+    secondary, accent, background, surface, and text colors.""",
+    backstory="""You are a world-renowned graphic designer with a deep understanding of 
+    color theory and its psychological impact. You create palettes that are not 
+    only beautiful but also functional and accessible (WCAG compliant).""",
+    verbose=True,
+    allow_delegation=False
+)
+
+# 2. Typography Specialist Agent
+typography_specialist = Agent(
+    role='Master Typographer for Digital Interfaces',
+    goal="""Define a complete and harmonious typography system, including font families 
+    and a responsive size scale (h1, h2, body, small).""",
+    backstory="""With decades of experience from print to digital, you excel at choosing 
+    fonts that are readable, stylish, and appropriate for the brand's voice. 
+    You create clear typographic hierarchies that work perfectly on the web.""",
+    verbose=True,
+    allow_delegation=False
+)
+
+# 3. Design System Architect Agent (The "Stitcher")
+design_system_architect = Agent(
+    role='Senior UI Design System Architect',
+    goal="""Consolidate the color and typography schemes into a single, comprehensive,
+    and developer-friendly JSON design system. Also add standard spacing and 
+    border-radius systems.""",
+    backstory="""You are a meticulous senior UI designer who creates robust and scalable 
+    design systems for large tech companies. Your work ensures absolute consistency 
+    and serves as the single source of truth for the entire development team.""",
+    verbose=True,
+    allow_delegation=False
+)
+
+# 4. Web UI Developer Agent
+web_developer = Agent(
+    role='Expert Frontend Developer specializing in Web UI',
+    goal="""Generate a complete, single HTML file for a given page description, 
+    strictly adhering to the provided design system. The code must be clean, 
+    responsive, and modern.""",
+    backstory="""You are a pixel-perfect frontend developer who translates design systems 
+    into clean, responsive, and maintainable HTML and CSS. You never deviate from 
+    the design specifications and your work is always of the highest quality.""",
+    verbose=True,
+    allow_delegation=False
+)
+
+
+# --- Task Definitions ---
+
+def create_design_tasks(prompt):
+    task_colors = Task(
+        description=f"""Analyze the user's prompt: '{prompt}'. 
+        Generate a JSON object representing the color palette. 
+        The JSON must contain keys: 'primary', 'secondary', 'accent', 
+        'background', 'surface', 'text', 'textLight', 'border'.""",
+        expected_output="A single, valid JSON object containing the color palette with 8 specified keys.",
+        agent=color_specialist
+    )
+
+    task_typography = Task(
+        description=f"""Analyze the user's prompt: '{prompt}'. 
+        Generate a JSON object for the typography system. 
+        It must include 'fontFamily' (web-safe), and sizes for 'h1', 'h2', 'body', and 'small'.""",
+        expected_output="A single, valid JSON object containing the typography system with 5 specified keys.",
+        agent=typography_specialist
+    )
+
+    task_architect = Task(
+        description="""Take the color palette and typography system from the specialists. 
+        Combine them into a single, final JSON object. Also, add standard 'spacing' 
+        (sm, md, lg, xl) and 'borderRadius' (sm, md, lg) systems. 
+        Your final output must be ONLY this complete JSON object and nothing else.""",
+        expected_output="A single, valid JSON object representing the complete design system.",
+        agent=design_system_architect,
+        context=[task_colors, task_typography]
+    )
+    return task_colors, task_typography, task_architect
+
+def create_developer_task(design_system_json, page_description):
+    return Task(
+        description=f"""Using the final design system provided below, create a complete HTML file 
+        for the following page: '{page_description}'. All CSS must be in a single `<style>` tag 
+        in the `<head>`.
+
+        FINAL DESIGN SYSTEM:
+        {design_system_json}""",
+        expected_output="The complete, raw HTML code for the webpage, starting with `<!DOCTYPE html>` and nothing else.",
+        agent=web_developer
+    )
+
+# --- Main Gradio Function ---
 
-# <<< NEW FUNCTION TO CREATE THE VISUAL STYLE GUIDE >>>
+def create_ui_design(prompt: str, api_key: str, progress=gr.Progress(track_tqdm=True)):
+    if not api_key:
+        return ("Missing API Key",) * 5
+    
+    # Configure the LLM to use OpenRouter
+    openrouter_llm = ChatOpenAI(
+        model=DEFAULT_MODEL,
+        api_key=api_key,
+        base_url="https://openrouter.ai/api/v1"
+    )
+
+    try:
+        # --- Part 1: Run the Design Crew to create the Design System ---
+        progress(0.1, desc="Briefing the design team...")
+        design_agents = [color_specialist, typography_specialist, design_system_architect]
+        design_tasks = create_design_tasks(prompt)
+        design_crew = Crew(
+            agents=design_agents,
+            tasks=list(design_tasks),
+            process=Process.sequential,
+            verbose=2,
+            llm=openrouter_llm
+        )
+        
+        progress(0.3, desc="Architecting the Design System...")
+        design_system_json_string = design_crew.kickoff(inputs={'prompt': prompt})
+        
+        # Parse the final JSON from the architect
+        design_system = json.loads(design_system_json_string)
+        
+        # --- Part 2: Run the Developer Agent for each page ---
+        progress(0.6, desc="Briefing the Web Developer...")
+        page_descriptions = re.findall(r'\d+\)([^0-9]+?)(?=\d+\)|$)', prompt) or [prompt]
+        while len(page_descriptions) < 3:
+            page_descriptions.append(f"Page {len(page_descriptions)+1} placeholder")
+
+        generated_pages = []
+        developer_crew = Crew(agents=[web_developer], tasks=[], process=Process.sequential, llm=openrouter_llm)
+
+        for i, page_desc in enumerate(page_descriptions[:3]):
+            progress(0.6 + (i * 0.1), desc=f"Developer is building Page {i+1}...")
+            dev_task = create_developer_task(design_system_json_string, page_desc)
+            developer_crew.tasks = [dev_task]
+            html_output = developer_crew.kickoff()
+            generated_pages.append(html_output)
+            
+        progress(0.9, desc="Creating Visualizer...")
+        visualizer_html = visualize_design_system(design_system)
+
+        progress(1.0, desc="Done!")
+        return generated_pages[0], generated_pages[1], generated_pages[2], design_system_json_string, visualizer_html
+
+    except Exception as e:
+        error_html = f"An error occurred with the CrewAI agents: {e}"
+        print(error_html)
+        return (error_html,) * 5
+
+# --- Visualizer (from previous answer, it's perfect) ---
 def visualize_design_system(design_system: Dict) -> str:
-    """
-    Takes a design system dictionary and returns an HTML string
-    that visually represents it as a style guide.
-    """
+    # This function is the same as the previous answer. It works perfectly.
+    # [Paste the entire visualize_design_system function from the previous answer here]
     if not design_system or not isinstance(design_system, dict):
         return "<div class='error-box'>No valid design system to display.</div>"
-
-    # Safely get data with defaults
     colors = design_system.get('colors', {})
     typography = design_system.get('typography', {})
-    spacing = design_system.get('spacing', {})
-    border_radius = design_system.get('borderRadius', {})
-
-    # ---- HTML & CSS for the visualizer itself ----
+    # ... rest of the function ...
     html = """
-    <!DOCTYPE html>
-    <html lang="en">
-    <head>
-        <meta charset="UTF-8">
-        <meta name="viewport" content="width=device-width, initial-scale=1.0">
-        <title>Design System</title>
-        <style>
-            body { 
-                font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, Helvetica, Arial, sans-serif;
-                background-color: #f9fafb; 
-                color: #374151;
-                padding: 1.5rem;
-                line-height: 1.6;
-            }
-            .section { margin-bottom: 2.5rem; }
-            h2 { 
-                font-size: 1.5rem; font-weight: 600; color: #111827;
-                border-bottom: 1px solid #e5e7eb; padding-bottom: 0.5rem; margin-bottom: 1.5rem;
-            }
-            .swatch-container { display: flex; flex-wrap: wrap; gap: 1.5rem; }
-            .swatch { 
-                width: 120px; height: 120px; border-radius: 0.75rem;
-                display: flex; flex-direction: column; justify-content: flex-end;
-                padding: 0.75rem; font-size: 0.875rem;
-                box-shadow: 0 4px 6px -1px rgba(0,0,0,0.1), 0 2px 4px -2px rgba(0,0,0,0.1);
-            }
-            .swatch .name { font-weight: 600; }
-            .swatch .hex { opacity: 0.8; }
-            
-            .type-container .sample { padding: 1rem; background-color: #fff; border: 1px solid #e5e7eb; border-radius: 0.5rem; margin-bottom: 1rem; }
-            .type-container .label { font-size: 0.9rem; color: #6b7280; margin-bottom: 0.25rem; }
-            
-            .component-container { display: flex; flex-direction: column; gap: 2rem; background-color: #fff; padding: 2rem; border-radius: 0.75rem; border: 1px solid #e5e7eb;}
-            .error-box { padding: 1rem; background: #fee2e2; color: #b91c1c; border-radius: 0.5rem; }
-        </style>
-    </head>
-    <body>
+    <!DOCTYPE html><html lang="en"><head><meta charset="UTF-8"><meta name="viewport" content="width=device-width, initial-scale=1.0"><title>Design System</title><style>body{font-family:-apple-system,BlinkMacSystemFont,'Segoe UI',Roboto,Helvetica,Arial,sans-serif;background-color:#f9fafb;color:#374151;padding:1.5rem;line-height:1.6}.section{margin-bottom:2.5rem}h2{font-size:1.5rem;font-weight:600;color:#111827;border-bottom:1px solid #e5e7eb;padding-bottom:.5rem;margin-bottom:1.5rem}.swatch-container{display:flex;flex-wrap:wrap;gap:1.5rem}.swatch{width:120px;height:120px;border-radius:.75rem;display:flex;flex-direction:column;justify-content:flex-end;padding:.75rem;font-size:.875rem;box-shadow:0 4px 6px -1px rgba(0,0,0,.1),0 2px 4px -2px rgba(0,0,0,.1)}.swatch .name{font-weight:600}.swatch .hex{opacity:.8}.type-container .sample{padding:1rem;background-color:#fff;border:1px solid #e5e7eb;border-radius:.5rem;margin-bottom:1rem}.type-container .label{font-size:.9rem;color:#6b7280;margin-bottom:.25rem}.component-container{display:flex;flex-direction:column;gap:2rem;background-color:#fff;padding:2rem;border-radius:.75rem;border:1px solid #e5e7eb}.error-box{padding:1rem;background:#fee2e2;color:#b91c1c;border-radius:.5rem}</style></head><body>
     """
-
-    # ---- 1. Color Palette Section ----
     html += "<div class='section'><h2>Color Palette</h2><div class='swatch-container'>"
     for name, hex_val in colors.items():
-        # Determine if text should be light or dark based on background
-        try:
-            r, g, b = int(hex_val[1:3], 16), int(hex_val[3:5], 16), int(hex_val[5:7], 16)
-            luminance = (0.299 * r + 0.587 * g + 0.114 * b) / 255
-            text_color = '#ffffff' if luminance < 0.5 else '#000000'
-        except:
-            text_color = '#000000'
-
-        html += f"""
-        <div class="swatch" style="background-color: {hex_val}; color: {text_color};">
-            <span class="name">{name.capitalize()}</span>
-            <span class="hex">{hex_val}</span>
-        </div>
-        """
+        try: r,g,b=int(hex_val[1:3],16),int(hex_val[3:5],16),int(hex_val[5:7],16); text_color='#ffffff' if (0.299*r+0.587*g+0.114*b)/255<0.5 else '#000000'
+        except: text_color='#000000'
+        html+=f'<div class="swatch" style="background-color:{hex_val};color:{text_color};"><span class="name">{name.capitalize()}</span><span class="hex">{hex_val}</span></div>'
     html += "</div></div>"
-    
-    # ---- 2. Typography Section ----
     font_family = typography.get('fontFamily', 'sans-serif')
-    html += f"<div class='section type-container'><h2>Typography</h2><p style='margin-bottom: 1.5rem;'>Font Family: <code>{font_family}</code></p>"
-    type_samples = {'h1': 'Heading 1', 'h2': 'Heading 2', 'body': 'This is a paragraph of body text. It is the default font size for long-form content.', 'small': 'Small utility text'}
-    for name, sample_text in type_samples.items():
-        font_size = typography.get(name, '1rem')
-        html += f"""
-        <div class="sample">
-            <div class="label">{name.capitalize()} ({font_size})</div>
-            <div style="font-family: {font_family}; font-size: {font_size};">
-                {sample_text}
-            </div>
-        </div>
-        """
+    html += f"<div class='section type-container'><h2>Typography</h2><p style='margin-bottom:1.5rem;'>Font Family: <code>{font_family}</code></p>"
+    for name, sample_text in {'h1':'Heading 1','h2':'Heading 2','body':'This is a paragraph of body text.','small':'Small utility text'}.items():
+        html+=f'<div class="sample"><div class="label">{name.capitalize()} ({typography.get(name,"1rem")})</div><div style="font-family:{font_family};font-size:{typography.get(name,"1rem")};">{sample_text}</div></div>'
     html += "</div>"
-
-    # ---- 3. Components Section ----
-    primary_color = colors.get('primary', '#3b82f6')
-    surface_color = colors.get('surface', '#f9fafb')
-    text_color = colors.get('text', '#111827')
-    border_color = colors.get('border', '#e5e7eb')
-    radius = border_radius.get('lg', '0.75rem')
-
-    html += f"""
-    <div class='section'>
-        <h2>Components Preview</h2>
-        <div class="component-container" style="background-color: {surface_color};">
-            <div>
-                <h3 style="font-family:{font_family}; font-size:{typography.get('h3', '1.25rem')}; color:{text_color}; margin-bottom: 1rem;">Example Card</h3>
-                <div style="background-color: #fff; border: 1px solid {border_color}; border-radius: {radius}; padding: 1.5rem; box-shadow: 0 1px 3px 0 rgba(0,0,0,0.1);">
-                    <p style="font-family:{font_family}; font-size:{typography.get('body', '1rem')}; color:{text_color}; margin-bottom: 1.5rem;">This is a sample card using the design system's surface, border, and radius values.</p>
-                    <button style="font-family:{font_family}; background-color: {primary_color}; color: white; border: none; padding: 0.75rem 1.5rem; border-radius: {border_radius.get('md', '0.5rem')}; cursor: pointer; font-size: 1rem;">Primary Button</button>
-                </div>
-            </div>
-        </div>
-    </div>
-    """
-
     html += "</body></html>"
     return html
 
-
-class UIDesignAgent:
-    # ... (The class code remains mostly the same, just ensuring it returns a proper dictionary) ...
-    def __init__(self, api_key: str):
-        self.api_key = api_key
-        self.base_url = "https://openrouter.ai/api/v1/chat/completions"
-        self.model = SELECTED_MODEL
-        self.design_system = {}
-
-    def get_default_design_system(self) -> Dict:
-        return {
-            "colors": { "primary": "#2563eb", "secondary": "#7c3aed", "accent": "#06b6d4", "background": "#ffffff", "surface": "#f8fafc", "text": "#1e293b", "textLight": "#64748b", "border": "#e2e8f0" },
-            "typography": { "fontFamily": "'Inter', system-ui, sans-serif", "h1": "2.5rem", "h2": "2rem", "h3": "1.5rem", "body": "1rem", "small": "0.875rem" },
-            "spacing": { "md": "1rem", "lg": "1.5rem", "xl": "2rem" },
-            "borderRadius": { "sm": "0.25rem", "md": "0.5rem", "lg": "0.75rem" }
-        }
-
-    def extract_design_system(self, prompt: str) -> Dict:
-        system_prompt = """You are an expert UI/UX designer. Create a consistent design system based on the user's prompt.
-        The design system must include: 'colors' (primary, secondary, background, surface, text, border), 
-        'typography' (fontFamily, h1, h2, body, small), and 'borderRadius' (sm, md, lg).
-        Return ONLY a single, valid JSON object and nothing else."""
-        headers = { "Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json" }
-        data = { "model": self.model, "messages": [{"role": "system", "content": system_prompt}, {"role": "user", "content": f"Create a design system for: {prompt}"}], "response_format": {"type": "json_object"}}
-        try:
-            response = requests.post(self.base_url, headers=headers, json=data, timeout=30)
-            result = response.json()
-            content = result['choices'][0]['message']['content']
-            self.design_system = json.loads(content)
-            return self.design_system
-        except Exception as e:
-            print(f"Error extracting design system: {e}. Using default.")
-            return self.get_default_design_system()
-
-    def generate_page_html(self, page_description: str) -> str:
-        # Use the stored design system
-        design_system_str = json.dumps(self.design_system)
-        system_prompt = f"""You are an expert UI developer. Create a complete, single HTML file.
-        Use this exact JSON design system for all styling: {design_system_str}
-        RULES:
-        1. Create a modern, responsive page.
-        2. Embed all CSS in a `<style>` tag in the `<head>`.
-        3. Your entire response must be ONLY raw HTML code, starting with `<!DOCTYPE html>`.
-        """
-        headers = { "Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json" }
-        data = { "model": self.model, "messages": [{"role": "system", "content": system_prompt}, {"role": "user", "content": f"Design page: {page_description}."}], "max_tokens": 3500 }
-        try:
-            response = requests.post(self.base_url, headers=headers, json=data, timeout=60)
-            return response.json()['choices'][0]['message']['content']
-        except Exception as e:
-            return f"<html><body>Error generating page: {e}</body></html>"
-
-    def generate_multi_page_design(self, prompt: str) -> Tuple[List[str], Dict]:
-        self.design_system = self.extract_design_system(prompt)
-        pages_desc = re.findall(r'\d+\)([^0-9]+?)(?=\d+\)|$)', prompt) or [prompt]
-        while len(pages_desc) < 3: pages_desc.append(f"Page {len(pages_desc)+1} placeholder")
-        generated_pages = [self.generate_page_html(desc) for desc in pages_desc[:3]]
-        return generated_pages, self.design_system
-
-def create_ui_design(prompt: str, api_key: str, progress=gr.Progress(track_tqdm=True)):
-    if not api_key: return ("Missing API Key",) * 5
-    
-    agent = UIDesignAgent(api_key)
-    progress(0.2, desc="Generating Design System...")
-    pages, design_system = agent.generate_multi_page_design(prompt)
-    
-    progress(0.8, desc="Creating Visualizer...")
-    # <<< CALL THE NEW VISUALIZER FUNCTION >>>
-    visualizer_html = visualize_design_system(design_system)
-    design_system_json_string = json.dumps(design_system, indent=2)
-
-    progress(1.0, desc="Done!")
-    return pages[0], pages[1], pages[2], design_system_json_string, visualizer_html
-
+# --- Gradio Interface ---
 def create_interface():
-    with gr.Blocks(title="AI UI Designer", theme=gr.themes.Soft()) as demo:
-        gr.Markdown("# 🎨 AI UI Designer - Consistent Multi-Page Generator")
+    # This UI layout is the same as the previous answer.
+    with gr.Blocks(title="CrewAI UI Designer", theme=gr.themes.Soft()) as demo:
+        gr.Markdown("# 🤖 CrewAI Powered UI Designer")
         with gr.Row():
             with gr.Column(scale=1):
                 api_key_input = gr.Textbox(label="OpenRouter API Key", placeholder="sk-or-...", type="password")
-                prompt_input = gr.Textbox(label="Website Description", placeholder="e.g., A 3-page site for a coffee shop: 1) landing page, 2) menu, 3) contact us", lines=4)
-                generate_btn = gr.Button("🚀 Generate Design", variant="primary")
-                
-                # <<< UPDATED UI WITH TABS FOR JSON and VISUALIZER >>>
+                prompt_input = gr.Textbox(label="Website Description", placeholder="e.g., A dark-mode, futuristic site for a synthwave musician: 1) landing page, 2) music, 3) tour dates", lines=4)
+                generate_btn = gr.Button("🚀 Assemble the Crew!", variant="primary")
                 with gr.Tabs():
                     with gr.TabItem("Visualizer"):
                         design_system_visualizer = gr.HTML(label="Design System Visualizer")
                     with gr.TabItem("JSON"):
                         design_system_output = gr.Code(label="Design System JSON", language="json", lines=15)
-
             with gr.Column(scale=2):
                 with gr.Tabs():
                     with gr.TabItem("📄 Page 1"): page1_output = gr.HTML()
                     with gr.TabItem("📄 Page 2"): page2_output = gr.HTML()
                     with gr.TabItem("📄 Page 3"): page3_output = gr.HTML()
-        
-        gr.Examples(
-            examples=["A dark-mode site for a DJ: 1) homepage with event list, 2) music gallery, 3) booking form", "A bright, clean site for a bakery: 1) landing page with featured items, 2) full menu page, 3) online order form"],
-            inputs=prompt_input
-        )
-        
-        # <<< UPDATE THE OUTPUTS TO INCLUDE THE NEW VISUALIZER COMPONENT >>>
         generate_btn.click(
             fn=create_ui_design, 
             inputs=[prompt_input, api_key_input],