Update app.py

app.py

@@ -10,46 +10,62 @@ class UIDesignAgent:
         self.api_key = api_key
         self.base_url = "https://openrouter.ai/api/v1/chat/completions"
         self.design_system = {}
-        self.color_palette = {}
-        self.typography = {}
         
     def extract_design_system(self, prompt: str) -> Dict:
         """Extract and maintain consistent design system across pages"""
         
-        system_prompt = """You are an expert UI/UX designer. Extract and define a consistent design system including:
-        1. Color Palette (primary, secondary, accent, neutral colors)
-        2. Typography (font families, sizes, weights)
-        3. Spacing system
-        4. Component styles
+        system_prompt = """You are an expert UI/UX designer. Create a consistent design system with:
+        - Primary color: #2563eb
+        - Secondary color: #7c3aed  
+        - Accent color: #06b6d4
+        - Background: #ffffff
+        - Text: #1e293b
+        - Font: Inter, system-ui, sans-serif
         
-        Return as JSON format."""
+        Use these EXACT colors for ALL pages. Return the design system as JSON."""
         
         headers = {
             "Authorization": f"Bearer {self.api_key}",
-            "Content-Type": "application/json"
+            "Content-Type": "application/json",
+            "HTTP-Referer": "https://huggingface.co",
+            "X-Title": "AI UI Designer"
         }
         
         data = {
-            "model": "google/gemini-pro",
+            "model": "openai/gpt-3.5-turbo",  # Using a more reliable model
             "messages": [
                 {"role": "system", "content": system_prompt},
                 {"role": "user", "content": f"Create a design system for: {prompt}"}
-            ]
+            ],
+            "temperature": 0.7,
+            "max_tokens": 1000
         }
         
         try:
             response = requests.post(self.base_url, headers=headers, json=data)
-            result = response.json()
-            design_system_text = result['choices'][0]['message']['content']
-            
-            # Extract JSON from response
-            json_match = re.search(r'\{.*\}', design_system_text, re.DOTALL)
-            if json_match:
-                self.design_system = json.loads(json_match.group())
-            
+            if response.status_code == 200:
+                result = response.json()
+                if 'choices' in result and len(result['choices']) > 0:
+                    design_system_text = result['choices'][0]['message']['content']
+                    
+                    # Try to extract JSON
+                    json_match = re.search(r'\{.*\}', design_system_text, re.DOTALL)
+                    if json_match:
+                        try:
+                            self.design_system = json.loads(json_match.group())
+                        except:
+                            self.design_system = self.get_default_design_system()
+                    else:
+                        self.design_system = self.get_default_design_system()
+                else:
+                    self.design_system = self.get_default_design_system()
+            else:
+                print(f"API Error: {response.status_code}")
+                self.design_system = self.get_default_design_system()
+                
             return self.design_system
         except Exception as e:
-            print(f"Error extracting design system: {e}")
+            print(f"Error: {e}")
             return self.get_default_design_system()
     
     def get_default_design_system(self) -> Dict:
@@ -62,90 +78,110 @@ class UIDesignAgent:
                 "background": "#ffffff",
                 "surface": "#f8fafc",
                 "text": "#1e293b",
-                "textLight": "#64748b"
+                "textLight": "#64748b",
+                "border": "#e2e8f0"
             },
             "typography": {
-                "fontFamily": "'Inter', -apple-system, sans-serif",
-                "h1": "48px",
-                "h2": "36px",
-                "h3": "24px",
-                "body": "16px",
-                "small": "14px"
+                "fontFamily": "'Inter', system-ui, -apple-system, sans-serif",
+                "h1": "2.5rem",
+                "h2": "2rem",
+                "h3": "1.5rem",
+                "body": "1rem",
+                "small": "0.875rem"
             },
             "spacing": {
-                "xs": "4px",
-                "sm": "8px",
-                "md": "16px",
-                "lg": "24px",
-                "xl": "32px"
+                "xs": "0.25rem",
+                "sm": "0.5rem",
+                "md": "1rem",
+                "lg": "1.5rem",
+                "xl": "2rem",
+                "xxl": "3rem"
+            },
+            "borderRadius": {
+                "sm": "0.25rem",
+                "md": "0.5rem",
+                "lg": "0.75rem",
+                "xl": "1rem"
             }
         }
     
     def generate_page_html(self, page_description: str, page_number: int) -> str:
         """Generate HTML for a specific page using consistent design system"""
         
-        design_context = f"""
-        Use this exact design system for consistency:
-        Colors: {json.dumps(self.design_system.get('colors', {}))}
-        Typography: {json.dumps(self.design_system.get('typography', {}))}
-        Spacing: {json.dumps(self.design_system.get('spacing', {}))}
-        """
+        colors = self.design_system.get('colors', self.get_default_design_system()['colors'])
+        typography = self.design_system.get('typography', self.get_default_design_system()['typography'])
         
-        system_prompt = f"""You are an expert UI developer. Create a complete HTML page with inline CSS.
-        {design_context}
+        system_prompt = f"""You are an expert UI developer. Create a complete HTML page.
         
-        Rules:
-        1. Use ONLY the provided color values
-        2. Use ONLY the provided typography settings
-        3. Create responsive, modern design
-        4. Include all CSS inline
-        5. Make it interactive with hover effects
-        6. Ensure consistency with the design system"""
+        MANDATORY DESIGN RULES:
+        - Primary color: {colors['primary']}
+        - Secondary color: {colors['secondary']}
+        - Background: {colors['background']}
+        - Text color: {colors['text']}
+        - Font family: {typography['fontFamily']}
+        
+        Use ONLY these colors. DO NOT use any other colors.
+        Create a modern, responsive design with these exact colors.
+        Include all CSS inline in the HTML.
+        """
         
         headers = {
             "Authorization": f"Bearer {self.api_key}",
-            "Content-Type": "application/json"
+            "Content-Type": "application/json",
+            "HTTP-Referer": "https://huggingface.co",
+            "X-Title": "AI UI Designer"
         }
         
         data = {
-            "model": "google/gemini-pro",
+            "model": "openai/gpt-3.5-turbo",
             "messages": [
                 {"role": "system", "content": system_prompt},
-                {"role": "user", "content": f"Create page {page_number}: {page_description}. Return ONLY the HTML code."}
-            ]
+                {"role": "user", "content": f"Create page {page_number}: {page_description}. Use the EXACT colors provided. Return complete HTML with inline CSS."}
+            ],
+            "temperature": 0.7,
+            "max_tokens": 2000
         }
         
         try:
             response = requests.post(self.base_url, headers=headers, json=data)
-            result = response.json()
-            html_content = result['choices'][0]['message']['content']
-            
-            # Extract HTML from code blocks if present
-            html_match = re.search(r'```html\n(.*?)\n```', html_content, re.DOTALL)
-            if html_match:
-                html_content = html_match.group(1)
-            
-            # Inject design system CSS
-            html_content = self.inject_design_system_css(html_content)
-            
-            return html_content
+            if response.status_code == 200:
+                result = response.json()
+                if 'choices' in result and len(result['choices']) > 0:
+                    html_content = result['choices'][0]['message']['content']
+                    
+                    # Extract HTML from code blocks
+                    html_match = re.search(r'```html?\n?(.*?)\n?```', html_content, re.DOTALL)
+                    if html_match:
+                        html_content = html_match.group(1)
+                    
+                    # Inject consistent CSS
+                    html_content = self.inject_design_system_css(html_content)
+                    return html_content
+                else:
+                    return self.create_error_page(f"No response for page {page_number}")
+            else:
+                return self.create_error_page(f"API Error: {response.status_code}")
         except Exception as e:
-            return f"<div style='color: red;'>Error generating page: {str(e)}</div>"
+            return self.create_error_page(f"Error: {str(e)}")
     
     def inject_design_system_css(self, html: str) -> str:
         """Inject consistent CSS variables into HTML"""
         
-        css_variables = f"""
+        colors = self.design_system.get('colors', self.get_default_design_system()['colors'])
+        typography = self.design_system.get('typography', self.get_default_design_system()['typography'])
+        
+        css_injection = f"""
         <style>
             :root {{
-                --primary: {self.design_system.get('colors', {}).get('primary', '#2563eb')};
-                --secondary: {self.design_system.get('colors', {}).get('secondary', '#7c3aed')};
-                --accent: {self.design_system.get('colors', {}).get('accent', '#06b6d4')};
-                --background: {self.design_system.get('colors', {}).get('background', '#ffffff')};
-                --surface: {self.design_system.get('colors', {}).get('surface', '#f8fafc')};
-                --text: {self.design_system.get('colors', {}).get('text', '#1e293b')};
-                --text-light: {self.design_system.get('colors', {}).get('textLight', '#64748b')};
-                --font-family: {self.design_system.get('typography', {}).get('fontFamily', 'Inter, sans-serif')};
+                --primary: {colors['primary']};
+                --secondary: {colors['secondary']};
+                --accent: {colors['accent']};
+                --background: {colors['background']};
+                --surface: {colors['surface']};
+                --text: {colors['text']};
+                --text-light: {colors['textLight']};
+                --border: {colors['border']};
+                --font-family: {typography['fontFamily']};
             }}
             
             * {{
@@ -161,175 +197,261 @@ class UIDesignAgent:
                 line-height: 1.6;
             }}
             
-            .primary-btn {{
+            h1, h2, h3, h4, h5, h6 {{
+                color: var(--text);
+                margin-bottom: 1rem;
+            }}
+            
+            a {{
+                color: var(--primary);
+                text-decoration: none;
+            }}
+            
+            a:hover {{
+                text-decoration: underline;
+            }}
+            
+            button, .btn {{
                 background-color: var(--primary);
                 color: white;
-                padding: 12px 24px;
+                padding: 0.75rem 1.5rem;
                 border: none;
-                border-radius: 8px;
+                border-radius: 0.5rem;
                 cursor: pointer;
+                font-family: var(--font-family);
                 transition: all 0.3s ease;
             }}
             
-            .primary-btn:hover {{
+            button:hover, .btn:hover {{
                 transform: translateY(-2px);
                 box-shadow: 0 4px 12px rgba(37, 99, 235, 0.3);
             }}
         </style>
         """
         
-        # Insert CSS after <head> tag or at the beginning
+        # Add viewport meta tag if not present
+        viewport_meta = '<meta name="viewport" content="width=device-width, initial-scale=1.0">'
+        
         if '<head>' in html:
-            html = html.replace('<head>', f'<head>{css_variables}')
+            html = html.replace('<head>', f'<head>\n{viewport_meta}\n{css_injection}')
+        elif '<html>' in html:
+            html = html.replace('<html>', f'<html>\n<head>\n{viewport_meta}\n{css_injection}\n</head>')
         else:
-            html = css_variables + html
-            
+            html = f"""<!DOCTYPE html>
+<html>
+<head>
+{viewport_meta}
+{css_injection}
+</head>
+<body>
+{html}
+</body>
+</html>"""
+        
         return html
     
+    def create_error_page(self, error_message: str) -> str:
+        """Create an error page with consistent styling"""
+        return f"""
+        <!DOCTYPE html>
+        <html>
+        <head>
+            <style>
+                body {{
+                    font-family: system-ui, -apple-system, sans-serif;
+                    display: flex;
+                    align-items: center;
+                    justify-content: center;
+                    height: 100vh;
+                    margin: 0;
+                    background: #f8fafc;
+                }}
+                .error {{
+                    background: white;
+                    padding: 2rem;
+                    border-radius: 0.5rem;
+                    box-shadow: 0 1px 3px rgba(0,0,0,0.1);
+                    max-width: 500px;
+                }}
+                h2 {{
+                    color: #ef4444;
+                    margin-bottom: 1rem;
+                }}
+                p {{
+                    color: #64748b;
+                }}
+            </style>
+        </head>
+        <body>
+            <div class="error">
+                <h2>Error Generating Page</h2>
+                <p>{error_message}</p>
+            </div>
+        </body>
+        </html>
+        """
+    
+    def extract_pages_from_prompt(self, prompt: str) -> List[str]:
+        """Extract individual page descriptions from prompt"""
+        
+        # Simple extraction based on common patterns
+        pages = []
+        
+        # Check for numbered pages
+        numbered_pattern = r'\d+\)([^0-9]+?)(?=\d+\)|$)'
+        matches = re.findall(numbered_pattern, prompt)
+        
+        if matches:
+            pages = [match.strip() for match in matches]
+        else:
+            # Check for comma-separated pages
+            if ',' in prompt:
+                pages = [p.strip() for p in prompt.split(',')]
+            else:
+                # Treat as single page
+                pages = [prompt]
+        
+        # Ensure we have at least 3 pages
+        while len(pages) < 3:
+            pages.append(f"Additional page {len(pages) + 1}")
+        
+        return pages[:3]  # Limit to 3 pages
+    
     def generate_multi_page_design(self, prompt: str) -> Tuple[List[str], Dict]:
         """Generate multiple pages with consistent design"""
         
-        # First, extract pages from prompt
-        pages = self.extract_pages_from_prompt(prompt)
-        
-        # Extract and set design system
+        # Extract design system first
         self.design_system = self.extract_design_system(prompt)
         
+        # Extract pages from prompt
+        pages = self.extract_pages_from_prompt(prompt)
+        
         # Generate HTML for each page
         generated_pages = []
-        for i, page_desc in enumerate(pages, 1):
+        for i, page_desc in enumerate(pages[:3], 1):  # Limit to 3 pages
             html = self.generate_page_html(page_desc, i)
             generated_pages.append(html)
         
-        return generated_pages, self.design_system
-    
-    def extract_pages_from_prompt(self, prompt: str) -> List[str]:
-        """Extract individual page descriptions from prompt"""
+        # Ensure we have exactly 3 pages
+        while len(generated_pages) < 3:
+            generated_pages.append(self.create_error_page("Page not specified"))
         
-        system_prompt = """Extract individual page descriptions from the user's request.
-        Return as a JSON array of page descriptions."""
-        
-        headers = {
-            "Authorization": f"Bearer {self.api_key}",
-            "Content-Type": "application/json"
-        }
-        
-        data = {
-            "model": "google/gemini-pro",
-            "messages": [
-                {"role": "system", "content": system_prompt},
-                {"role": "user", "content": prompt}
-            ]
-        }
-        
-        try:
-            response = requests.post(self.base_url, headers=headers, json=data)
-            result = response.json()
-            pages_text = result['choices'][0]['message']['content']
-            
-            # Extract JSON array
-            json_match = re.search(r'\[.*\]', pages_text, re.DOTALL)
-            if json_match:
-                pages = json.loads(json_match.group())
-                return pages
-            else:
-                # Default to single page if extraction fails
-                return [prompt]
-        except Exception as e:
-            print(f"Error extracting pages: {e}")
-            return [prompt]
+        return generated_pages, self.design_system
 
-# Gradio Interface
 def create_ui_design(prompt: str, api_key: str):
     """Main function to create UI designs"""
     
-    if not api_key:
-        return ["<div style='color: red;'>Please provide OpenRouter API key</div>"], "No design system"
-    
-    agent = UIDesignAgent(api_key)
-    pages, design_system = agent.generate_multi_page_design(prompt)
-    
-    # Format design system for display
-    design_system_display = json.dumps(design_system, indent=2)
+    if not api_key or api_key.strip() == "":
+        error_html = """
+        <div style='padding: 20px; color: #ef4444; background: #fee2e2; border-radius: 8px;'>
+            <h3>API Key Required</h3>
+            <p>Please provide your OpenRouter API key to generate designs.</p>
+            <p>Get your key from <a href='https://openrouter.ai' target='_blank'>OpenRouter.ai</a></p>
+        </div>
+        """
+        return error_html, error_html, error_html, "{}"
     
-    # Ensure we have at least 3 pages for the interface
-    while len(pages) < 3:
-        pages.append("<div style='padding: 20px;'>Page placeholder</div>")
+    if not prompt or prompt.strip() == "":
+        error_html = """
+        <div style='padding: 20px; color: #f59e0b; background: #fef3c7; border-radius: 8px;'>
+            <h3>Prompt Required</h3>
+            <p>Please describe the website you want to create.</p>
+        </div>
+        """
+        return error_html, error_html, error_html, "{}"
     
-    return pages[0], pages[1] if len(pages) > 1 else "", pages[2] if len(pages) > 2 else "", design_system_display
+    try:
+        agent = UIDesignAgent(api_key)
+        pages, design_system = agent.generate_multi_page_design(prompt)
+        
+        # Format design system for display
+        design_system_display = json.dumps(design_system, indent=2)
+        
+        return pages[0], pages[1], pages[2], design_system_display
+        
+    except Exception as e:
+        error_html = f"""
+        <div style='padding: 20px; color: #ef4444; background: #fee2e2; border-radius: 8px;'>
+            <h3>Generation Error</h3>
+            <p>{str(e)}</p>
+        </div>
+        """
+        return error_html, error_html, error_html, "{}"
 
 # Create Gradio interface
-with gr.Blocks(title="AI UI Designer - Stitch Clone", theme=gr.themes.Soft()) as demo:
-    gr.Markdown("""
-    # 🎨 AI UI Designer (Stitch-like)
-    ### Generate consistent multi-page website designs with AI
-    
-    This tool ensures consistent colors, typography, and styling across all generated pages.
-    """)
-    
-    with gr.Row():
-        with gr.Column(scale=1):
-            api_key_input = gr.Textbox(
-                label="OpenRouter API Key",
-                placeholder="Enter your OpenRouter API key",
-                type="password"
-            )
-            
-            prompt_input = gr.Textbox(
-                label="Design Prompt",
-                placeholder="Describe your multi-page website (e.g., 'Create a 3-page website for a tech startup with: 1) Landing page with hero section, 2) About page with team section, 3) Contact page with form')",
-                lines=5
-            )
-            
-            generate_btn = gr.Button("🚀 Generate Design", variant="primary")
-            
-            design_system_output = gr.Code(
-                label="Design System",
-                language="json",
-                lines=10
-            )
-        
-        with gr.Column(scale=2):
-            with gr.Tabs():
-                with gr.TabItem("Page 1"):
-                    page1_output = gr.HTML(label="Page 1 Preview")
+def create_interface():
+    with gr.Blocks(title="AI UI Designer - Consistent Multi-Page", theme=gr.themes.Soft()) as demo:
+        gr.Markdown("""
+        # 🎨 AI UI Designer - Consistent Multi-Page Generator
+        ### Generate multiple website pages with consistent design system
+        
+        This tool ensures **consistent colors and styling** across all generated pages.
+        """)
+        
+        with gr.Row():
+            with gr.Column(scale=1):
+                api_key_input = gr.Textbox(
+                    label="OpenRouter API Key",
+                    placeholder="sk-or-...",
+                    type="password",
+                    info="Get your key from openrouter.ai"
+                )
                 
-                with gr.TabItem("Page 2"):
-                    page2_output = gr.HTML(label="Page 2 Preview")
+                prompt_input = gr.Textbox(
+                    label="Website Description",
+                    placeholder="Example: Create a tech startup website with: 1) Landing page with hero section and features, 2) About page with team section, 3) Contact page with form",
+                    lines=4
+                )
                 
-                with gr.TabItem("Page 3"):
-                    page3_output = gr.HTML(label="Page 3 Preview")
-    
-    gr.Markdown("""
-    ### Features:
-    - ✅ Consistent color scheme across all pages
-    - ✅ Unified typography and spacing
-    - ✅ Responsive design
-    - ✅ Modern UI components
-    - ✅ Design system extraction
-    
-    ### How to use:
-    1. Get your API key from [OpenRouter](https://openrouter.ai/)
-    2. Describe your multi-page website
-    3. Click Generate Design
-    4. View consistent designs across all pages
-    """)
-    
-    generate_btn.click(
-        fn=create_ui_design,
-        inputs=[prompt_input, api_key_input],
-        outputs=[page1_output, page2_output, page3_output, design_system_output]
-    )
-    
-    # Example prompts
-    gr.Examples(
-        examples=[
-            ["Create a 3-page SaaS website: 1) Landing page with pricing cards, 2) Features page with grid layout, 3) Contact page with form"],
-            ["Design a portfolio website: 1) Home with hero section, 2) Projects gallery page, 3) About me page with skills"],
-            ["Build an e-commerce site: 1) Homepage with product grid, 2) Product detail page, 3) Shopping cart page"],
-        ],
-        inputs=prompt_input
-    )
-
-if __name__ == "__main__":
-    demo.launch()
+                with gr.Row():
+                    generate_btn = gr.Button("🚀 Generate Design", variant="primary", scale=2)
+                    clear_btn = gr.Button("🗑️ Clear", scale=1)
+                
+                gr.Markdown("### 📋 Design System")
+                design_system_output = gr.Code(
+                    language="json",
+                    lines=10
+                )
+            
+            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.Markdown("""
+        ---
+        ### 💡 Tips for Best Results:
+        - Clearly describe each page you want
+        - Specify key components (hero, navbar, footer, etc.)
+        - Mention any specific features you need
+        
+        ### 🎯 Features:
+        - ✅ Consistent color scheme across all pages
+        - ✅ Unified typography and spacing
+        - ✅ Modern, responsive design
+        - ✅ Professional UI components
+        """)
+        
+        # Set up event handlers
+        generate_btn.click(
+            fn=create_ui_design,
+            inputs=[prompt_input, api_key_input],
+            outputs=[page1_output, page2_output, page3_output, design_system_output]
+        )
+        
+        clear_btn.click(
+            fn=lambda: ("", "", "", "", "{}"),
+            outputs=[prompt_input, page1_output, page2_output, page3_output, design_system_output]
+        )
+        
+        # Add examples
+        gr.Examples(
+            examples=[
+                "Create a SaaS website: 1) Landing page with hero and pricing cards, 2) Features page with grid layout, 3) Contact page with form",
+                "Design a portfolio: 1) Homepage with projects showcase, 2) About page with