Upload app.py

app.py

@@ -849,6 +849,734 @@ def normalized_to_dict(normalized) -> dict:
     return result
 
 
+# =============================================================================
+# STAGE 2: NEW ARCHITECTURE (Rule Engine + Benchmark Research + LLM Agents)
+# =============================================================================
+
+async def run_stage2_analysis_v2(
+    selected_benchmarks: list[str] = None,
+    progress=gr.Progress()
+):
+    """
+    Run Stage 2 analysis with new architecture:
+    - Layer 1: Rule Engine (FREE)
+    - Layer 2: Benchmark Research (Firecrawl + Cache)
+    - Layer 3: LLM Agents (Brand ID, Benchmark Advisor, Best Practices)
+    - Layer 4: HEAD Synthesizer
+    
+    Includes comprehensive error handling for graceful degradation.
+    """
+    
+    # Validate Stage 1 completion
+    if not state.desktop_normalized or not state.mobile_normalized:
+        return create_stage2_error_response("❌ Please complete Stage 1 first")
+    
+    # Default benchmarks if none selected
+    if not selected_benchmarks or len(selected_benchmarks) == 0:
+        selected_benchmarks = [
+            "material_design_3",
+            "shopify_polaris", 
+            "atlassian_design",
+        ]
+    
+    state.log("")
+    state.log("═" * 60)
+    state.log("🚀 STAGE 2: MULTI-AGENT ANALYSIS")
+    state.log("═" * 60)
+    state.log(f"   Started: {datetime.now().strftime('%H:%M:%S')}")
+    state.log(f"   Benchmarks: {', '.join(selected_benchmarks)}")
+    state.log("")
+    
+    # Initialize results with defaults (for graceful degradation)
+    rule_results = None
+    benchmark_comparisons = []
+    brand_result = None
+    benchmark_advice = None
+    best_practices = None
+    final_synthesis = None
+    
+    progress(0.05, desc="⚙️ Running Rule Engine...")
+    
+    try:
+        # =================================================================
+        # LAYER 1: RULE ENGINE (FREE) - Critical, must succeed
+        # =================================================================
+        try:
+            from core.rule_engine import run_rule_engine
+            
+            # Convert tokens to dict
+            desktop_dict = normalized_to_dict(state.desktop_normalized)
+            mobile_dict = normalized_to_dict(state.mobile_normalized)
+            
+            # Validate we have data
+            if not desktop_dict.get("colors") and not desktop_dict.get("typography"):
+                raise ValueError("No tokens extracted from Stage 1")
+            
+            # Run rule engine
+            rule_results = run_rule_engine(
+                typography_tokens=desktop_dict.get("typography", {}),
+                color_tokens=desktop_dict.get("colors", {}),
+                spacing_tokens=desktop_dict.get("spacing", {}),
+                radius_tokens=desktop_dict.get("radius", {}),
+                shadow_tokens=desktop_dict.get("shadows", {}),
+                log_callback=state.log,
+            )
+            
+            state.rule_engine_results = rule_results
+            state.log("")
+            state.log("   ✅ Rule Engine: SUCCESS")
+            
+        except Exception as e:
+            state.log(f"   ❌ Rule Engine FAILED: {str(e)[:100]}")
+            state.log("   └─ Cannot proceed without rule engine results")
+            import traceback
+            state.log(traceback.format_exc()[:500])
+            return create_stage2_error_response(f"❌ Rule Engine failed: {str(e)}")
+        
+        progress(0.20, desc="🔬 Researching benchmarks...")
+        
+        # =================================================================
+        # LAYER 2: BENCHMARK RESEARCH - Can use fallback
+        # =================================================================
+        try:
+            from agents.benchmark_researcher import BenchmarkResearcher, FALLBACK_BENCHMARKS, BenchmarkData
+            
+            # Try to get Firecrawl client (optional)
+            firecrawl_client = None
+            try:
+                from agents.firecrawl_extractor import get_firecrawl_client
+                firecrawl_client = get_firecrawl_client()
+                state.log("   ├─ Firecrawl client: Available")
+            except Exception as fc_err:
+                state.log(f"   ├─ Firecrawl client: Not available ({str(fc_err)[:30]})")
+                state.log("   │  └─ Will use cached/fallback data")
+            
+            # Get HF client for LLM extraction (optional)
+            hf_client = None
+            try:
+                from core.hf_inference import get_inference_client
+                hf_client = get_inference_client()
+                state.log("   ├─ HF client: Available")
+            except Exception as hf_err:
+                state.log(f"   ├─ HF client: Not available ({str(hf_err)[:30]})")
+            
+            researcher = BenchmarkResearcher(
+                firecrawl_client=firecrawl_client,
+                hf_client=hf_client,
+            )
+            
+            # Research selected benchmarks (with fallback)
+            try:
+                benchmarks = await researcher.research_selected_benchmarks(
+                    selected_keys=selected_benchmarks,
+                    log_callback=state.log,
+                )
+            except Exception as research_err:
+                state.log(f"   ⚠️ Research failed, using fallback: {str(research_err)[:50]}")
+                # Use fallback data
+                benchmarks = []
+                for key in selected_benchmarks:
+                    if key in FALLBACK_BENCHMARKS:
+                        data = FALLBACK_BENCHMARKS[key]
+                        benchmarks.append(BenchmarkData(
+                            key=key,
+                            name=key.replace("_", " ").title(),
+                            short_name=key.split("_")[0].title(),
+                            vendor="",
+                            icon="📦",
+                            typography=data.get("typography", {}),
+                            spacing=data.get("spacing", {}),
+                            colors=data.get("colors", {}),
+                            fetched_at=datetime.now().isoformat(),
+                            confidence="fallback",
+                            best_for=[],
+                        ))
+            
+            # Compare to benchmarks
+            if benchmarks and rule_results:
+                benchmark_comparisons = researcher.compare_to_benchmarks(
+                    your_ratio=rule_results.typography.detected_ratio,
+                    your_base_size=int(rule_results.typography.sizes_px[0]) if rule_results.typography.sizes_px else 16,
+                    your_spacing_grid=rule_results.spacing.detected_base,
+                    benchmarks=benchmarks,
+                    log_callback=state.log,
+                )
+                state.benchmark_comparisons = benchmark_comparisons
+                state.log("")
+                state.log(f"   ✅ Benchmark Research: SUCCESS ({len(benchmarks)} systems)")
+            else:
+                state.log("   ⚠️ No benchmarks available for comparison")
+            
+        except Exception as e:
+            state.log(f"   ⚠️ Benchmark Research FAILED: {str(e)[:100]}")
+            state.log("   └─ Continuing without benchmark comparison...")
+            benchmark_comparisons = []
+        
+        progress(0.40, desc="🤖 Running LLM Agents...")
+        
+        # =================================================================
+        # LAYER 3: LLM AGENTS - Can fail gracefully
+        # =================================================================
+        try:
+            from agents.llm_agents import (
+                BrandIdentifierAgent,
+                BenchmarkAdvisorAgent,
+                BestPracticesValidatorAgent,
+                BrandIdentification,
+                BenchmarkAdvice,
+                BestPracticesResult,
+            )
+            
+            state.log("")
+            state.log("═" * 60)
+            state.log("🤖 LAYER 3: LLM ANALYSIS")
+            state.log("═" * 60)
+            
+            # Check if HF client is available
+            if not hf_client:
+                try:
+                    from core.hf_inference import get_inference_client
+                    hf_client = get_inference_client()
+                except Exception:
+                    state.log("   ⚠️ HF client not available - skipping LLM agents")
+                    hf_client = None
+            
+            if hf_client:
+                # Initialize agents
+                brand_agent = BrandIdentifierAgent(hf_client)
+                benchmark_agent = BenchmarkAdvisorAgent(hf_client)
+                best_practices_agent = BestPracticesValidatorAgent(hf_client)
+                
+                # Get semantic analysis from Stage 1
+                semantic_analysis = getattr(state, 'semantic_analysis', {})
+                desktop_dict = normalized_to_dict(state.desktop_normalized)
+                
+                # Run agents (with individual error handling)
+                # Brand Identifier
+                try:
+                    brand_result = await brand_agent.analyze(
+                        color_tokens=desktop_dict.get("colors", {}),
+                        semantic_analysis=semantic_analysis,
+                        log_callback=state.log,
+                    )
+                except Exception as e:
+                    state.log(f"   ⚠️ Brand Identifier failed: {str(e)[:50]}")
+                    brand_result = BrandIdentification()
+                
+                # Benchmark Advisor
+                if benchmark_comparisons:
+                    try:
+                        benchmark_advice = await benchmark_agent.analyze(
+                            user_ratio=rule_results.typography.detected_ratio,
+                            user_base=int(rule_results.typography.sizes_px[0]) if rule_results.typography.sizes_px else 16,
+                            user_spacing=rule_results.spacing.detected_base,
+                            benchmark_comparisons=benchmark_comparisons,
+                            log_callback=state.log,
+                        )
+                    except Exception as e:
+                        state.log(f"   ⚠️ Benchmark Advisor failed: {str(e)[:50]}")
+                        benchmark_advice = BenchmarkAdvice()
+                else:
+                    benchmark_advice = BenchmarkAdvice()
+                
+                # Best Practices Validator
+                try:
+                    best_practices = await best_practices_agent.analyze(
+                        rule_engine_results=rule_results,
+                        log_callback=state.log,
+                    )
+                except Exception as e:
+                    state.log(f"   ⚠️ Best Practices Validator failed: {str(e)[:50]}")
+                    best_practices = BestPracticesResult(overall_score=rule_results.consistency_score)
+            else:
+                # No HF client - use defaults
+                state.log("   └─ Using default values (no LLM)")
+                brand_result = BrandIdentification()
+                benchmark_advice = BenchmarkAdvice()
+                best_practices = BestPracticesResult(overall_score=rule_results.consistency_score)
+                
+        except Exception as e:
+            state.log(f"   ⚠️ LLM Agents FAILED: {str(e)[:100]}")
+            brand_result = BrandIdentification() if not brand_result else brand_result
+            benchmark_advice = BenchmarkAdvice() if not benchmark_advice else benchmark_advice
+            best_practices = BestPracticesResult(overall_score=rule_results.consistency_score if rule_results else 50)
+        
+        progress(0.70, desc="🧠 Synthesizing results...")
+        
+        # =================================================================
+        # LAYER 4: HEAD SYNTHESIZER - Can use fallback
+        # =================================================================
+        try:
+            from agents.llm_agents import HeadSynthesizerAgent, HeadSynthesis
+            
+            if hf_client and brand_result and benchmark_advice and best_practices:
+                head_agent = HeadSynthesizerAgent(hf_client)
+                
+                try:
+                    final_synthesis = await head_agent.synthesize(
+                        rule_engine_results=rule_results,
+                        benchmark_comparisons=benchmark_comparisons,
+                        brand_identification=brand_result,
+                        benchmark_advice=benchmark_advice,
+                        best_practices=best_practices,
+                        log_callback=state.log,
+                    )
+                except Exception as e:
+                    state.log(f"   ⚠️ HEAD Synthesizer failed: {str(e)[:50]}")
+                    final_synthesis = None
+            
+            # Create fallback synthesis if needed
+            if not final_synthesis:
+                state.log("   └─ Creating fallback synthesis...")
+                final_synthesis = create_fallback_synthesis(
+                    rule_results, benchmark_comparisons, brand_result, best_practices
+                )
+            
+            state.final_synthesis = final_synthesis
+            
+        except Exception as e:
+            state.log(f"   ⚠️ Synthesis FAILED: {str(e)[:100]}")
+            final_synthesis = create_fallback_synthesis(
+                rule_results, benchmark_comparisons, brand_result, best_practices
+            )
+            state.final_synthesis = final_synthesis
+        
+        progress(0.85, desc="📊 Formatting results...")
+        
+        # =================================================================
+        # FORMAT OUTPUTS FOR UI
+        # =================================================================
+        
+        try:
+            # Build status markdown
+            status_md = format_stage2_status_v2(
+                rule_results=rule_results,
+                final_synthesis=final_synthesis,
+                best_practices=best_practices,
+            )
+            
+            # Build benchmark comparison HTML
+            benchmark_md = format_benchmark_comparison_v2(
+                benchmark_comparisons=benchmark_comparisons,
+                benchmark_advice=benchmark_advice,
+            )
+            
+            # Build scores dashboard HTML
+            scores_html = format_scores_dashboard_v2(
+                rule_results=rule_results,
+                final_synthesis=final_synthesis,
+                best_practices=best_practices,
+            )
+            
+            # Build priority actions HTML
+            actions_html = format_priority_actions_v2(
+                rule_results=rule_results,
+                final_synthesis=final_synthesis,
+                best_practices=best_practices,
+            )
+            
+            # Build color recommendations table
+            color_recs_table = format_color_recommendations_table_v2(
+                rule_results=rule_results,
+                brand_result=brand_result,
+                final_synthesis=final_synthesis,
+            )
+            
+            # Get fonts and typography data
+            fonts = get_detected_fonts()
+            base_size = get_base_font_size()
+            
+            typography_desktop_data = format_typography_comparison_viewport(
+                state.desktop_normalized, base_size, "desktop"
+            )
+            typography_mobile_data = format_typography_comparison_viewport(
+                state.mobile_normalized, base_size, "mobile"
+            )
+            
+            # Generate visual previews
+            typography_preview_html = ""
+            try:
+                from core.preview_generator import generate_typography_preview_html
+                
+                primary_font = fonts.get("primary", "Open Sans")
+                desktop_typo_dict = {
+                    name: {
+                        "font_size": t.font_size,
+                        "font_weight": t.font_weight,
+                        "line_height": t.line_height,
+                    }
+                    for name, t in state.desktop_normalized.typography.items()
+                }
+                typography_preview_html = generate_typography_preview_html(desktop_typo_dict, primary_font)
+            except Exception as preview_err:
+                state.log(f"   ⚠️ Preview generation failed: {str(preview_err)[:50]}")
+                typography_preview_html = "<div style='padding:20px;background:#f5f5f5;'>Preview unavailable</div>"
+            
+        except Exception as format_err:
+            state.log(f"   ⚠️ Formatting failed: {str(format_err)[:100]}")
+            # Return minimal results
+            return (
+                f"⚠️ Analysis completed with formatting errors: {str(format_err)[:50]}",
+                state.get_logs(),
+                "*Benchmark comparison unavailable*",
+                "<div>Scores unavailable</div>",
+                "<div>Actions unavailable</div>",
+                [],
+                None,
+                None,
+                "",
+            )
+        
+        progress(0.95, desc="✅ Complete!")
+        
+        # Final log summary
+        state.log("")
+        state.log("═" * 60)
+        state.log("📊 FINAL RESULTS")
+        state.log("═" * 60)
+        state.log("")
+        overall_score = final_synthesis.scores.get('overall', rule_results.consistency_score) if final_synthesis else rule_results.consistency_score
+        state.log(f"   🎯 OVERALL SCORE: {overall_score}/100")
+        if final_synthesis and final_synthesis.scores:
+            state.log(f"   ├─ Accessibility:  {final_synthesis.scores.get('accessibility', '?')}/100")
+            state.log(f"   ├─ Consistency:    {final_synthesis.scores.get('consistency', '?')}/100")
+            state.log(f"   └─ Organization:   {final_synthesis.scores.get('organization', '?')}/100")
+        state.log("")
+        if benchmark_comparisons:
+            state.log(f"   🏆 Closest Benchmark: {benchmark_comparisons[0].benchmark.name if benchmark_comparisons else 'N/A'}")
+        state.log("")
+        state.log("   🎯 TOP 3 ACTIONS:")
+        if final_synthesis and final_synthesis.top_3_actions:
+            for i, action in enumerate(final_synthesis.top_3_actions[:3]):
+                impact = action.get('impact', 'medium')
+                icon = "🔴" if impact == "high" else "🟡" if impact == "medium" else "🟢"
+                state.log(f"   │  {i+1}. {icon} {action.get('action', 'N/A')}")
+        else:
+            state.log(f"   │  1. 🔴 Fix {rule_results.aa_failures} AA compliance failures")
+        state.log("")
+        state.log("═" * 60)
+        state.log(f"   💰 TOTAL COST: ~$0.003")
+        state.log(f"   ⏱️  COMPLETED: {datetime.now().strftime('%H:%M:%S')}")
+        state.log("═" * 60)
+        
+        return (
+            status_md,
+            state.get_logs(),
+            benchmark_md,
+            scores_html,
+            actions_html,
+            color_recs_table,
+            typography_desktop_data,
+            typography_mobile_data,
+            typography_preview_html,
+        )
+        
+    except Exception as e:
+        import traceback
+        state.log(f"❌ Critical Error: {str(e)}")
+        state.log(traceback.format_exc())
+        return create_stage2_error_response(f"❌ Analysis failed: {str(e)}")
+
+
+def create_fallback_synthesis(rule_results, benchmark_comparisons, brand_result, best_practices):
+    """Create a fallback synthesis when LLM synthesis fails."""
+    from agents.llm_agents import HeadSynthesis
+    
+    # Calculate scores from rule engine
+    overall = rule_results.consistency_score if rule_results else 50
+    accessibility = max(0, 100 - (rule_results.aa_failures * 10)) if rule_results else 50
+    
+    # Build actions from rule engine
+    actions = []
+    if rule_results and rule_results.aa_failures > 0:
+        actions.append({
+            "action": f"Fix {rule_results.aa_failures} colors failing AA compliance",
+            "impact": "high",
+            "effort": "30 min",
+        })
+    if rule_results and not rule_results.typography.is_consistent:
+        actions.append({
+            "action": f"Align type scale to {rule_results.typography.recommendation} ({rule_results.typography.recommendation_name})",
+            "impact": "medium",
+            "effort": "1 hour",
+        })
+    if rule_results and rule_results.color_stats.unique_count > 30:
+        actions.append({
+            "action": f"Consolidate {rule_results.color_stats.unique_count} colors to ~15 semantic colors",
+            "impact": "medium",
+            "effort": "2 hours",
+        })
+    
+    return HeadSynthesis(
+        executive_summary=f"Your design system scores {overall}/100. Analysis completed with fallback synthesis.",
+        scores={
+            "overall": overall,
+            "accessibility": accessibility,
+            "consistency": overall,
+            "organization": 50,
+        },
+        benchmark_fit={
+            "closest": benchmark_comparisons[0].benchmark.name if benchmark_comparisons else "Unknown",
+            "similarity": f"{benchmark_comparisons[0].overall_match_pct:.0f}%" if benchmark_comparisons else "N/A",
+        },
+        brand_analysis={
+            "primary": brand_result.brand_primary.get("color", "Unknown") if brand_result else "Unknown",
+            "cohesion": brand_result.cohesion_score if brand_result else 5,
+        },
+        top_3_actions=actions[:3],
+        color_recommendations=[],
+        type_scale_recommendation={
+            "current_ratio": rule_results.typography.detected_ratio if rule_results else 1.0,
+            "recommended_ratio": rule_results.typography.recommendation if rule_results else 1.25,
+        },
+        spacing_recommendation={
+            "current": f"{rule_results.spacing.detected_base}px" if rule_results else "Unknown",
+            "recommended": f"{rule_results.spacing.recommendation}px" if rule_results else "8px",
+        },
+    )
+
+
+def create_stage2_error_response(error_msg: str):
+    """Create error response tuple for Stage 2."""
+    return (
+        error_msg,
+        state.get_logs(),
+        "",  # benchmark_md
+        "",  # scores_html
+        "",  # actions_html
+        [],  # color_recs_table
+        None,  # typography_desktop
+        None,  # typography_mobile
+        "",  # typography_preview
+    )
+
+
+def format_stage2_status_v2(rule_results, final_synthesis, best_practices) -> str:
+    """Format Stage 2 status with new architecture results."""
+    
+    lines = []
+    lines.append("## ✅ Analysis Complete!")
+    lines.append("")
+    
+    # Overall Score
+    overall = final_synthesis.scores.get('overall', rule_results.consistency_score)
+    lines.append(f"### 🎯 Overall Score: {overall}/100")
+    lines.append("")
+    
+    # Executive Summary
+    if final_synthesis.executive_summary:
+        lines.append(f"*{final_synthesis.executive_summary}*")
+        lines.append("")
+    
+    # Quick Stats
+    lines.append("### 📊 Quick Stats")
+    lines.append(f"- **AA Failures:** {rule_results.aa_failures}")
+    lines.append(f"- **Type Scale:** {rule_results.typography.detected_ratio:.3f} ({rule_results.typography.scale_name})")
+    lines.append(f"- **Spacing Grid:** {rule_results.spacing.detected_base}px ({rule_results.spacing.alignment_percentage:.0f}% aligned)")
+    lines.append(f"- **Unique Colors:** {rule_results.color_stats.unique_count}")
+    lines.append("")
+    
+    # Cost
+    lines.append("### 💰 Cost")
+    lines.append("**Total:** ~$0.003 (Rule Engine: $0 + LLM: ~$0.003)")
+    
+    return "\n".join(lines)
+
+
+def format_benchmark_comparison_v2(benchmark_comparisons, benchmark_advice) -> str:
+    """Format benchmark comparison results."""
+    
+    if not benchmark_comparisons:
+        return "*No benchmark comparison available*"
+    
+    lines = []
+    lines.append("## 📊 Benchmark Comparison")
+    lines.append("")
+    
+    # Recommended benchmark
+    if benchmark_advice and benchmark_advice.recommended_benchmark_name:
+        lines.append(f"### 🏆 Recommended: {benchmark_advice.recommended_benchmark_name}")
+        if benchmark_advice.reasoning:
+            lines.append(f"*{benchmark_advice.reasoning[:200]}*")
+        lines.append("")
+    
+    # Comparison table
+    lines.append("### 📈 Similarity Ranking")
+    lines.append("")
+    lines.append("| Rank | Design System | Match | Type Ratio | Base | Grid |")
+    lines.append("|------|---------------|-------|------------|------|------|")
+    
+    medals = ["🥇", "🥈", "🥉"]
+    for i, c in enumerate(benchmark_comparisons[:5]):
+        medal = medals[i] if i < 3 else str(i+1)
+        b = c.benchmark
+        lines.append(
+            f"| {medal} | {b.icon} {b.short_name} | {c.overall_match_pct:.0f}% | "
+            f"{b.typography.get('scale_ratio', '?')} | {b.typography.get('base_size', '?')}px | "
+            f"{b.spacing.get('base', '?')}px |"
+        )
+    
+    lines.append("")
+    
+    # Alignment changes needed
+    if benchmark_advice and benchmark_advice.alignment_changes:
+        lines.append("### 🔧 Changes to Align")
+        for change in benchmark_advice.alignment_changes[:3]:
+            lines.append(f"- **{change.get('change', '?')}**: {change.get('from', '?')} → {change.get('to', '?')} (effort: {change.get('effort', '?')})")
+    
+    return "\n".join(lines)
+
+
+def format_scores_dashboard_v2(rule_results, final_synthesis, best_practices) -> str:
+    """Format scores dashboard HTML."""
+    
+    overall = final_synthesis.scores.get('overall', rule_results.consistency_score)
+    accessibility = final_synthesis.scores.get('accessibility', 100 - (rule_results.aa_failures * 5))
+    consistency = final_synthesis.scores.get('consistency', rule_results.consistency_score)
+    organization = final_synthesis.scores.get('organization', 50)
+    
+    def score_color(score):
+        if score >= 80:
+            return "#10b981"  # Green
+        elif score >= 60:
+            return "#f59e0b"  # Yellow
+        else:
+            return "#ef4444"  # Red
+    
+    html = f"""
+    <div style="display: grid; grid-template-columns: repeat(4, 1fr); gap: 16px; margin: 20px 0;">
+        <div style="background: linear-gradient(135deg, {score_color(overall)}22 0%, {score_color(overall)}11 100%); 
+                    border: 2px solid {score_color(overall)}; border-radius: 12px; padding: 20px; text-align: center;">
+            <div style="font-size: 32px; font-weight: 700; color: {score_color(overall)};">{overall}</div>
+            <div style="font-size: 12px; color: #64748b; margin-top: 4px;">OVERALL</div>
+        </div>
+        <div style="background: #f8fafc; border: 1px solid #e2e8f0; border-radius: 12px; padding: 20px; text-align: center;">
+            <div style="font-size: 24px; font-weight: 600; color: {score_color(accessibility)};">{accessibility}</div>
+            <div style="font-size: 12px; color: #64748b; margin-top: 4px;">Accessibility</div>
+        </div>
+        <div style="background: #f8fafc; border: 1px solid #e2e8f0; border-radius: 12px; padding: 20px; text-align: center;">
+            <div style="font-size: 24px; font-weight: 600; color: {score_color(consistency)};">{consistency}</div>
+            <div style="font-size: 12px; color: #64748b; margin-top: 4px;">Consistency</div>
+        </div>
+        <div style="background: #f8fafc; border: 1px solid #e2e8f0; border-radius: 12px; padding: 20px; text-align: center;">
+            <div style="font-size: 24px; font-weight: 600; color: {score_color(organization)};">{organization}</div>
+            <div style="font-size: 12px; color: #64748b; margin-top: 4px;">Organization</div>
+        </div>
+    </div>
+    """
+    
+    return html
+
+
+def format_priority_actions_v2(rule_results, final_synthesis, best_practices) -> str:
+    """Format priority actions HTML."""
+    
+    actions = final_synthesis.top_3_actions if final_synthesis.top_3_actions else []
+    
+    # If no synthesis actions, build from rule engine
+    if not actions and best_practices and best_practices.priority_fixes:
+        actions = best_practices.priority_fixes
+    
+    if not actions:
+        # Default actions from rule engine
+        actions = []
+        if rule_results.aa_failures > 0:
+            actions.append({
+                "action": f"Fix {rule_results.aa_failures} colors failing AA compliance",
+                "impact": "high",
+                "effort": "30 min",
+            })
+        if not rule_results.typography.is_consistent:
+            actions.append({
+                "action": f"Align type scale to {rule_results.typography.recommendation} ({rule_results.typography.recommendation_name})",
+                "impact": "medium", 
+                "effort": "1 hour",
+            })
+        if rule_results.color_stats.unique_count > 30:
+            actions.append({
+                "action": f"Consolidate {rule_results.color_stats.unique_count} colors to ~15 semantic colors",
+                "impact": "medium",
+                "effort": "2 hours",
+            })
+    
+    html_items = []
+    for i, action in enumerate(actions[:3]):
+        impact = action.get('impact', 'medium')
+        border_color = "#ef4444" if impact == "high" else "#f59e0b" if impact == "medium" else "#10b981"
+        impact_bg = "#fee2e2" if impact == "high" else "#fef3c7" if impact == "medium" else "#dcfce7"
+        impact_text = "#991b1b" if impact == "high" else "#92400e" if impact == "medium" else "#166534"
+        icon = "🔴" if impact == "high" else "🟡" if impact == "medium" else "🟢"
+        
+        html_items.append(f"""
+        <div style="background: white; border: 1px solid #e2e8f0; border-left: 4px solid {border_color}; 
+                    border-radius: 8px; padding: 16px; margin-bottom: 12px;">
+            <div style="display: flex; justify-content: space-between; align-items: flex-start;">
+                <div>
+                    <div style="font-weight: 600; color: #1e293b; margin-bottom: 4px;">
+                        {icon} {action.get('action', 'N/A')}
+                    </div>
+                    <div style="font-size: 13px; color: #64748b;">
+                        {action.get('details', '')}
+                    </div>
+                </div>
+                <div style="display: flex; gap: 8px;">
+                    <span style="background: {impact_bg}; color: {impact_text}; padding: 4px 8px; 
+                                 border-radius: 12px; font-size: 11px; font-weight: 600;">
+                        {impact.upper()}
+                    </span>
+                    <span style="background: #f1f5f9; color: #475569; padding: 4px 8px; 
+                                 border-radius: 12px; font-size: 11px;">
+                        {action.get('effort', '?')}
+                    </span>
+                </div>
+            </div>
+        </div>
+        """)
+    
+    return f"""
+    <div style="margin: 20px 0;">
+        <h3 style="margin-bottom: 16px; color: #1e293b;">🎯 Priority Actions</h3>
+        {''.join(html_items)}
+    </div>
+    """
+
+
+def format_color_recommendations_table_v2(rule_results, brand_result, final_synthesis) -> list:
+    """Format color recommendations as table data."""
+    
+    rows = []
+    
+    # Add AA failures with fixes
+    for a in rule_results.accessibility:
+        if not a.passes_aa_normal and a.suggested_fix:
+            role = "brand.primary" if brand_result and brand_result.brand_primary.get("color") == a.hex_color else a.name
+            rows.append([
+                True,  # Accept checkbox
+                role,
+                a.hex_color,
+                f"Fails AA ({a.contrast_on_white:.1f}:1)",
+                a.suggested_fix,
+                f"{a.suggested_fix_contrast:.1f}:1",
+            ])
+    
+    # Add recommendations from synthesis
+    if final_synthesis and final_synthesis.color_recommendations:
+        for rec in final_synthesis.color_recommendations:
+            if rec.get("current") != rec.get("suggested"):
+                # Check if not already in rows
+                if not any(r[2] == rec.get("current") for r in rows):
+                    rows.append([
+                        rec.get("accept", True),
+                        rec.get("role", "unknown"),
+                        rec.get("current", ""),
+                        rec.get("reason", ""),
+                        rec.get("suggested", ""),
+                        "",
+                    ])
+    
+    return rows
+
+
 def build_analysis_status(final_recs: dict, cost_tracking: dict, errors: list) -> str:
     """Build status markdown from analysis results."""
     
@@ -2116,24 +2844,223 @@ def export_tokens_json():
 # =============================================================================
 
 def create_ui():
-    """Create the Gradio interface."""
+    """Create the Gradio interface with corporate branding."""
+    
+    # Corporate theme customization
+    corporate_theme = gr.themes.Base(
+        primary_hue=gr.themes.colors.blue,
+        secondary_hue=gr.themes.colors.slate,
+        neutral_hue=gr.themes.colors.slate,
+        font=[gr.themes.GoogleFont("Inter"), "ui-sans-serif", "system-ui", "sans-serif"],
+        font_mono=[gr.themes.GoogleFont("JetBrains Mono"), "ui-monospace", "monospace"],
+    ).set(
+        # Colors
+        body_background_fill="#f8fafc",
+        body_background_fill_dark="#0f172a",
+        block_background_fill="white",
+        block_background_fill_dark="#1e293b",
+        block_border_color="#e2e8f0",
+        block_border_color_dark="#334155",
+        block_label_background_fill="#f1f5f9",
+        block_label_background_fill_dark="#1e293b",
+        block_title_text_color="#0f172a",
+        block_title_text_color_dark="#f1f5f9",
+        
+        # Primary button
+        button_primary_background_fill="#2563eb",
+        button_primary_background_fill_hover="#1d4ed8",
+        button_primary_text_color="white",
+        
+        # Secondary button
+        button_secondary_background_fill="#f1f5f9",
+        button_secondary_background_fill_hover="#e2e8f0",
+        button_secondary_text_color="#1e293b",
+        
+        # Input fields
+        input_background_fill="#ffffff",
+        input_background_fill_dark="#1e293b",
+        input_border_color="#cbd5e1",
+        input_border_color_dark="#475569",
+        
+        # Shadows and radius
+        block_shadow="0 1px 3px rgba(0,0,0,0.1)",
+        block_shadow_dark="0 1px 3px rgba(0,0,0,0.3)",
+        block_border_width="1px",
+        block_radius="8px",
+        
+        # Text
+        body_text_color="#1e293b",
+        body_text_color_dark="#e2e8f0",
+        body_text_size="14px",
+    )
+    
+    # Custom CSS for additional styling
+    custom_css = """
+    /* Global styles */
+    .gradio-container {
+        max-width: 1400px !important;
+        margin: 0 auto !important;
+    }
+    
+    /* Header branding */
+    .app-header {
+        background: linear-gradient(135deg, #1e40af 0%, #3b82f6 100%);
+        padding: 24px 32px;
+        border-radius: 12px;
+        margin-bottom: 24px;
+        color: white;
+    }
+    .app-header h1 {
+        margin: 0 0 8px 0;
+        font-size: 28px;
+        font-weight: 700;
+    }
+    .app-header p {
+        margin: 0;
+        opacity: 0.9;
+        font-size: 14px;
+    }
+    
+    /* Stage indicators */
+    .stage-header {
+        background: linear-gradient(90deg, #f1f5f9 0%, #ffffff 100%);
+        padding: 16px 20px;
+        border-radius: 8px;
+        border-left: 4px solid #2563eb;
+        margin-bottom: 16px;
+    }
+    .stage-header h2 {
+        margin: 0;
+        font-size: 18px;
+        color: #1e293b;
+    }
+    
+    /* Log styling */
+    .log-container textarea {
+        font-family: 'JetBrains Mono', monospace !important;
+        font-size: 12px !important;
+        line-height: 1.6 !important;
+        background: #0f172a !important;
+        color: #e2e8f0 !important;
+        border-radius: 8px !important;
+    }
+    
+    /* Color swatch */
+    .color-swatch {
+        display: inline-block;
+        width: 24px;
+        height: 24px;
+        border-radius: 4px;
+        margin-right: 8px;
+        vertical-align: middle;
+        border: 1px solid rgba(0,0,0,0.1);
+    }
+    
+    /* Score badges */
+    .score-badge {
+        display: inline-block;
+        padding: 4px 12px;
+        border-radius: 20px;
+        font-weight: 600;
+        font-size: 13px;
+    }
+    .score-badge.high { background: #dcfce7; color: #166534; }
+    .score-badge.medium { background: #fef3c7; color: #92400e; }
+    .score-badge.low { background: #fee2e2; color: #991b1b; }
+    
+    /* Benchmark cards */
+    .benchmark-card {
+        background: #f8fafc;
+        border: 1px solid #e2e8f0;
+        border-radius: 8px;
+        padding: 16px;
+        margin-bottom: 12px;
+    }
+    .benchmark-card.selected {
+        border-color: #2563eb;
+        background: #eff6ff;
+    }
+    
+    /* Action items */
+    .action-item {
+        background: white;
+        border: 1px solid #e2e8f0;
+        border-radius: 8px;
+        padding: 16px;
+        margin-bottom: 8px;
+    }
+    .action-item.high-priority {
+        border-left: 4px solid #ef4444;
+    }
+    .action-item.medium-priority {
+        border-left: 4px solid #f59e0b;
+    }
+    
+    /* Progress indicator */
+    .progress-bar {
+        height: 4px;
+        background: #e2e8f0;
+        border-radius: 2px;
+        overflow: hidden;
+    }
+    .progress-bar-fill {
+        height: 100%;
+        background: linear-gradient(90deg, #2563eb, #3b82f6);
+        transition: width 0.3s ease;
+    }
+    
+    /* Accordion styling */
+    .accordion-header {
+        font-weight: 600 !important;
+    }
+    
+    /* Table styling */
+    table {
+        border-collapse: collapse;
+        width: 100%;
+    }
+    th {
+        background: #f1f5f9;
+        padding: 12px;
+        text-align: left;
+        font-weight: 600;
+        border-bottom: 2px solid #e2e8f0;
+    }
+    td {
+        padding: 12px;
+        border-bottom: 1px solid #e2e8f0;
+    }
+    
+    /* Dark mode adjustments */
+    .dark .stage-header {
+        background: linear-gradient(90deg, #1e293b 0%, #0f172a 100%);
+        border-left-color: #3b82f6;
+    }
+    .dark .stage-header h2 {
+        color: #f1f5f9;
+    }
+    .dark .benchmark-card {
+        background: #1e293b;
+        border-color: #334155;
+    }
+    .dark .action-item {
+        background: #1e293b;
+        border-color: #334155;
+    }
+    """
     
     with gr.Blocks(
         title="Design System Extractor v2",
-        theme=gr.themes.Soft(),
-        css="""
-        .color-swatch { display: inline-block; width: 24px; height: 24px; border-radius: 4px; margin-right: 8px; vertical-align: middle; }
-        """
+        theme=corporate_theme,
+        css=custom_css
     ) as app:
         
-        gr.Markdown("""
-        # 🎨 Design System Extractor v2
-        
-        **Reverse-engineer design systems from live websites.**
-        
-        A semi-automated, human-in-the-loop system that extracts, normalizes, and upgrades design tokens.
-        
-        ---
+        # Header with branding
+        gr.HTML("""
+        <div class="app-header">
+            <h1>🎨 Design System Extractor v2</h1>
+            <p>Reverse-engineer design systems from live websites • AI-powered analysis • Figma-ready export</p>
+        </div>
         """)
         
         # =================================================================
@@ -2294,60 +3221,138 @@ def create_ui():
         # STAGE 2: AI UPGRADES
         # =================================================================
         
-        with gr.Accordion("🧠 Stage 2: AI-Powered Upgrades", open=False) as stage2_accordion:
+        with gr.Accordion("🧠 Stage 2: AI-Powered Analysis", open=False) as stage2_accordion:
+            
+            # Stage header
+            gr.HTML("""
+            <div class="stage-header">
+                <h2>🧠 Stage 2: Multi-Agent Analysis</h2>
+                <p style="color: #64748b; margin-top: 4px;">Rule Engine + Benchmark Research + LLM Agents</p>
+            </div>
+            """)
             
             stage2_status = gr.Markdown("Click 'Analyze' to start AI-powered design system analysis.")
             
             # =============================================================
-            # LLM CONFIGURATION & COMPETITORS
+            # NEW ARCHITECTURE CONFIGURATION
             # =============================================================
-            with gr.Accordion("⚙️ Analysis Configuration", open=False):
+            with gr.Accordion("⚙️ Analysis Configuration", open=True):
+                
+                # Architecture explanation
                 gr.Markdown("""
-                ### 🤖 LLM Models Used
+                ### 🏗️ New Analysis Architecture
                 
-                | Role | Model | Expertise |
-                |------|-------|-----------|
-                | **Typography Analyst** | meta-llama/Llama-3.1-70B | Type scale patterns, readability |
-                | **Color Analyst** | meta-llama/Llama-3.1-70B | Color theory, accessibility |
-                | **Spacing Analyst** | Rule-based | Grid alignment, consistency |
+                | Layer | Type | What It Does | Cost |
+                |-------|------|--------------|------|
+                | **Layer 1** | Rule Engine | Type scale, AA check, spacing grid, color stats | FREE |
+                | **Layer 2** | Benchmark Research | Fetch live specs via Firecrawl (24h cache) | ~$0.001 |
+                | **Layer 3** | LLM Agents | Brand ID, Benchmark Advisor, Best Practices | ~$0.002 |
+                | **Layer 4** | HEAD Synthesizer | Combine all → Final recommendations | ~$0.001 |
                 
-                *Analysis compares your design against industry leaders.*
+                **Total Cost:** ~$0.003-0.004 per analysis
                 """)
                 
-                gr.Markdown("### 🎯 Competitor Design Systems")
-                gr.Markdown("Enter design systems to compare against (comma-separated):")
-                competitors_input = gr.Textbox(
-                    value="Material Design 3, Apple HIG, Shopify Polaris, IBM Carbon, Atlassian",
-                    label="Competitors",
-                    placeholder="Material Design 3, Apple HIG, Shopify Polaris...",
+                gr.Markdown("---")
+                
+                # Benchmark selection
+                gr.Markdown("### 📊 Select Design Systems to Compare Against")
+                gr.Markdown("*Choose which design systems to benchmark your tokens against:*")
+                
+                benchmark_checkboxes = gr.CheckboxGroup(
+                    choices=[
+                        ("🟢 Material Design 3 (Google)", "material_design_3"),
+                        ("🍎 Apple HIG", "apple_hig"),
+                        ("🛒 Shopify Polaris", "shopify_polaris"),
+                        ("🔵 Atlassian Design System", "atlassian_design"),
+                        ("🔷 IBM Carbon", "ibm_carbon"),
+                        ("🌊 Tailwind CSS", "tailwind_css"),
+                        ("🐜 Ant Design", "ant_design"),
+                        ("⚡ Chakra UI", "chakra_ui"),
+                    ],
+                    value=["material_design_3", "shopify_polaris", "atlassian_design"],
+                    label="Benchmarks",
+                )
+                
+                gr.Markdown("""
+                <small style="color: #64748b;">
+                💡 <b>Tip:</b> Select 2-4 benchmarks for best results. More benchmarks = longer analysis time.
+                <br>
+                📦 Results are cached for 24 hours to speed up subsequent analyses.
+                </small>
+                """)
+            
+            # Analyze button
+            with gr.Row():
+                analyze_btn_v2 = gr.Button(
+                    "🚀 Run Analysis (New Architecture)", 
+                    variant="primary", 
+                    size="lg",
+                    scale=2
+                )
+                analyze_btn_legacy = gr.Button(
+                    "🤖 Legacy Analysis", 
+                    variant="secondary",
+                    size="lg",
+                    scale=1
+                )
+            
+            # =============================================================
+            # ANALYSIS LOG
+            # =============================================================
+            with gr.Accordion("📋 Analysis Log", open=True):
+                stage2_log = gr.Textbox(
+                    label="Log", 
+                    lines=20, 
+                    interactive=False,
+                    elem_classes=["log-container"]
                 )
-                gr.Markdown("*Suggestions: Ant Design, Chakra UI, Tailwind, Bootstrap, Salesforce Lightning*")
             
-            analyze_btn = gr.Button("🤖 Analyze Design System", variant="primary", size="lg")
+            # =============================================================
+            # SCORES DASHBOARD
+            # =============================================================
+            gr.Markdown("---")
+            gr.Markdown("## 📊 Analysis Results")
+            
+            scores_dashboard = gr.HTML(
+                value="<div style='padding: 40px; background: #f8fafc; border-radius: 12px; text-align: center; color: #94a3b8;'>Scores will appear after analysis...</div>",
+                label="Scores"
+            )
             
-            with gr.Accordion("📋 AI Analysis Log", open=True):
-                stage2_log = gr.Textbox(label="Log", lines=18, interactive=False)
+            # =============================================================
+            # PRIORITY ACTIONS
+            # =============================================================
+            priority_actions_html = gr.HTML(
+                value="<div style='padding: 20px; background: #f8fafc; border-radius: 8px; color: #94a3b8;'>Priority actions will appear after analysis...</div>",
+                label="Priority Actions"
+            )
             
             # =============================================================
-            # BRAND COMPARISON (LLM Research)
+            # BENCHMARK COMPARISON
             # =============================================================
             gr.Markdown("---")
-            brand_comparison = gr.Markdown("*Brand comparison will appear after analysis*")
+            benchmark_comparison_md = gr.Markdown("*Benchmark comparison will appear after analysis*")
             
             # =============================================================
-            # FONT FAMILIES DETECTED
+            # COLOR RECOMMENDATIONS
             # =============================================================
             gr.Markdown("---")
-            gr.Markdown("## 🔤 Font Families Detected")
-            font_families_display = gr.Markdown("*Font information will appear after analysis*")
+            gr.Markdown("## 🎨 Color Recommendations")
+            gr.Markdown("*Accept or reject AI-suggested color changes:*")
+            
+            color_recommendations_table = gr.Dataframe(
+                headers=["Accept", "Role", "Current", "Issue", "Suggested", "New Contrast"],
+                datatype=["bool", "str", "str", "str", "str", "str"],
+                label="Color Recommendations",
+                interactive=True,
+                row_count=(0, "dynamic"),
+            )
             
             # =============================================================
-            # TYPOGRAPHY SECTION - Desktop & Mobile
+            # TYPOGRAPHY SECTION
             # =============================================================
             gr.Markdown("---")
             gr.Markdown("## 📐 Typography")
             
-            # Visual Preview
             with gr.Accordion("👁️ Typography Visual Preview", open=True):
                 stage2_typography_preview = gr.HTML(
                     value="<div style='padding: 20px; background: #f5f5f5; border-radius: 8px; color: #666;'>Typography preview will appear after analysis...</div>",
@@ -2541,11 +3546,28 @@ def create_ui():
             outputs=[colors_table, typography_table, spacing_table],
         )
         
-        # Stage 2: Analyze
-        analyze_btn.click(
+        # Stage 2: NEW Architecture Analyze
+        analyze_btn_v2.click(
+            fn=run_stage2_analysis_v2,
+            inputs=[benchmark_checkboxes],
+            outputs=[
+                stage2_status, 
+                stage2_log, 
+                benchmark_comparison_md,
+                scores_dashboard,
+                priority_actions_html,
+                color_recommendations_table,
+                typography_desktop, 
+                typography_mobile, 
+                stage2_typography_preview,
+            ],
+        )
+        
+        # Stage 2: Legacy Analyze (keep for backward compatibility)
+        analyze_btn_legacy.click(
             fn=run_stage2_analysis,
-            inputs=[competitors_input],
-            outputs=[stage2_status, stage2_log, brand_comparison, font_families_display,
+            inputs=[],
+            outputs=[stage2_status, stage2_log, benchmark_comparison_md, scores_dashboard,
                      typography_desktop, typography_mobile, spacing_comparison, 
                      base_colors_display, color_ramps_display, radius_display, shadows_display,
                      stage2_typography_preview, stage2_color_ramps_preview,
@@ -2577,9 +3599,11 @@ def create_ui():
         
         gr.Markdown("""
         ---
-        **Design System Extractor v2** | Built with Playwright + Gradio + LangGraph + HuggingFace
+        **Design System Extractor v2** | Built with Playwright + Firecrawl + LangGraph + HuggingFace
         
         *A semi-automated co-pilot for design system recovery and modernization.*
+        
+        **New Architecture:** Rule Engine (FREE) + Benchmark Research (Firecrawl) + LLM Agents
         """)
     
     return app