Upload stage2_graph.py

agents/stage2_graph.py

@@ -53,6 +53,7 @@ class Stage2State(TypedDict):
     desktop_tokens: dict
     mobile_tokens: dict
     competitors: list[str]
+    semantic_analysis: Optional[dict]  # NEW: Semantic color categorization from Stage 1
     
     # Parallel analysis outputs
     llm1_analysis: Optional[dict]
@@ -302,13 +303,15 @@ async def analyze_with_llm1(state: Stage2State, log_callback: Optional[Callable]
         log_callback(f"   Provider: {provider}")
         log_callback(f"   💰 Cost: ${llm1_config.get('cost_per_million_input', 0.29)}/M in, ${llm1_config.get('cost_per_million_output', 0.59)}/M out")
         log_callback(f"   📝 Task: Typography, Colors, AA, Spacing analysis")
+        log_callback(f"   🧠 Semantic context: {'Yes' if state.get('semantic_analysis') else 'No'}")
         log_callback("")
     
-    # Build prompt
+    # Build prompt with semantic analysis
     prompt = build_analyst_prompt(
         tokens_summary=summarize_tokens(state["desktop_tokens"], state["mobile_tokens"]),
         competitors=state["competitors"],
         persona=llm1_config.get("persona", "Senior Design Systems Architect"),
+        semantic_analysis=state.get("semantic_analysis"),
     )
     
     try:
@@ -408,13 +411,15 @@ async def analyze_with_llm2(state: Stage2State, log_callback: Optional[Callable]
         log_callback(f"   Provider: {provider}")
         log_callback(f"   💰 Cost: ${llm2_config.get('cost_per_million_input', 0.59)}/M in, ${llm2_config.get('cost_per_million_output', 0.79)}/M out")
         log_callback(f"   📝 Task: Typography, Colors, AA, Spacing analysis")
+        log_callback(f"   🧠 Semantic context: {'Yes' if state.get('semantic_analysis') else 'No'}")
         log_callback("")
     
-    # Build prompt
+    # Build prompt with semantic analysis
     prompt = build_analyst_prompt(
         tokens_summary=summarize_tokens(state["desktop_tokens"], state["mobile_tokens"]),
         competitors=state["competitors"],
         persona=llm2_config.get("persona", "Senior Design Systems Architect"),
+        semantic_analysis=state.get("semantic_analysis"),
     )
     
     try:
@@ -567,13 +572,14 @@ async def compile_with_head(state: Stage2State, log_callback: Optional[Callable]
         log_callback(f"   Provider: {provider}")
         log_callback(f"   💰 Cost: ${head_config.get('cost_per_million_input', 0.59)}/M in, ${head_config.get('cost_per_million_output', 0.79)}/M out")
         log_callback("")
-        log_callback("   📥 INPUT: Analyzing outputs from LLM 1 + LLM 2 + Rules...")
+        log_callback("   📥 INPUT: Analyzing outputs from LLM 1 + LLM 2 + Rules + Semantic...")
     
-    # Build HEAD prompt
+    # Build HEAD prompt with semantic context
     prompt = build_head_prompt(
         llm1_analysis=state.get("llm1_analysis", {}),
         llm2_analysis=state.get("llm2_analysis", {}),
         rule_calculations=state.get("rule_calculations", {}),
+        semantic_analysis=state.get("semantic_analysis"),
     )
     
     try:
@@ -707,8 +713,60 @@ def summarize_tokens(desktop: dict, mobile: dict) -> str:
     return "\n".join(lines)
 
 
-def build_analyst_prompt(tokens_summary: str, competitors: list[str], persona: str) -> str:
-    """Build prompt for analyst LLMs."""
+def build_analyst_prompt(tokens_summary: str, competitors: list[str], persona: str, semantic_analysis: dict = None) -> str:
+    """Build prompt for analyst LLMs with semantic color context."""
+    
+    # Build semantic colors section if available
+    semantic_section = ""
+    brand_primary_hex = "unknown"
+    
+    if semantic_analysis:
+        semantic_section = "\n\n## SEMANTIC COLOR ANALYSIS (from Stage 1)\n"
+        semantic_section += "Colors have been categorized by their actual CSS usage:\n"
+        
+        # Brand colors
+        brand = semantic_analysis.get("brand", {})
+        if brand:
+            semantic_section += "\n### 🎨 Brand Colors (used on buttons, CTAs, links)\n"
+            for role, data in brand.items():
+                if data and isinstance(data, dict):
+                    hex_val = data.get('hex', 'N/A')
+                    if role == "primary":
+                        brand_primary_hex = hex_val
+                    semantic_section += f"- **{role}**: {hex_val} ({data.get('confidence', 'unknown')} confidence)\n"
+        
+        # Text colors
+        text = semantic_analysis.get("text", {})
+        if text:
+            semantic_section += "\n### 📝 Text Colors (used with 'color' CSS property)\n"
+            for role, data in text.items():
+                if data and isinstance(data, dict):
+                    semantic_section += f"- **{role}**: {data.get('hex', 'N/A')}\n"
+        
+        # Background colors
+        bg = semantic_analysis.get("background", {})
+        if bg:
+            semantic_section += "\n### 🖼️ Background Colors (used with 'background-color')\n"
+            for role, data in bg.items():
+                if data and isinstance(data, dict):
+                    semantic_section += f"- **{role}**: {data.get('hex', 'N/A')}\n"
+        
+        # Border colors
+        border = semantic_analysis.get("border", {})
+        if border:
+            semantic_section += "\n### 📏 Border Colors\n"
+            for role, data in border.items():
+                if data and isinstance(data, dict):
+                    semantic_section += f"- **{role}**: {data.get('hex', 'N/A')}\n"
+        
+        # Feedback colors
+        feedback = semantic_analysis.get("feedback", {})
+        if feedback:
+            semantic_section += "\n### 🚨 Feedback Colors\n"
+            for role, data in feedback.items():
+                if data and isinstance(data, dict):
+                    semantic_section += f"- **{role}**: {data.get('hex', 'N/A')}\n"
+
     return f"""You are a {persona}.
 
 ## YOUR TASK
@@ -716,6 +774,7 @@ Analyze these design tokens extracted from a website and compare against industr
 
 ## EXTRACTED TOKENS
 {tokens_summary}
+{semantic_section}
 
 ## COMPETITOR DESIGN SYSTEMS TO RESEARCH
 {', '.join(competitors)}
@@ -728,14 +787,18 @@ Analyze these design tokens extracted from a website and compare against industr
 - Compare to competitors: what ratios do they use?
 - Score (1-10) and specific recommendations
 
-### 2. Colors
-- Is the color palette cohesive?
-- Are semantic colors properly defined (primary, secondary, etc.)?
-- Score (1-10) and specific recommendations
+### 2. Colors (USE SEMANTIC ANALYSIS ABOVE!)
+- **IMPORTANT**: The brand primary color is {brand_primary_hex}
+- Is this brand color appropriate? Check contrast on white/light backgrounds
+- Does the text color hierarchy (primary → secondary → muted) work well?
+- Are feedback colors (error, success, warning) properly defined?
+- Score (1-10) and SPECIFIC recommendations per color role
 
 ### 3. Accessibility (AA Compliance)
-- What contrast issues might exist?
-- Score (1-10)
+- Check brand primary color ({brand_primary_hex}) contrast on white
+- Check text primary color contrast
+- Identify any colors that fail WCAG AA (4.5:1 for text, 3:1 for large text)
+- Score (1-10) and specific fixes with suggested replacement hex values
 
 ### 4. Spacing
 - Is spacing consistent? Does it follow a grid (4px, 8px)?
@@ -747,9 +810,29 @@ Analyze these design tokens extracted from a website and compare against industr
 ## RESPOND IN JSON FORMAT ONLY:
 ```json
 {{
-  "typography": {{"analysis": "...", "detected_ratio": 1.2, "score": 7, "recommendations": ["..."]}},
-  "colors": {{"analysis": "...", "score": 6, "recommendations": ["..."]}},
-  "accessibility": {{"issues": ["..."], "score": 5}},
+  "typography": {{
+    "analysis": "...", 
+    "detected_ratio": 1.2, 
+    "score": 7, 
+    "recommendations": ["..."]
+  }},
+  "colors": {{
+    "analysis": "...", 
+    "brand_primary": "{brand_primary_hex}",
+    "brand_primary_issue": "...",
+    "brand_primary_suggestion": "#xxx or keep",
+    "score": 6, 
+    "role_recommendations": [
+      {{"role": "brand.primary", "current": "#xxx", "action": "keep|darken|lighten", "suggested": "#xxx", "reason": "..."}},
+      {{"role": "text.primary", "current": "#xxx", "action": "keep|darken|lighten", "suggested": "#xxx", "reason": "..."}}
+    ]
+  }},
+  "accessibility": {{
+    "issues": [
+      {{"color": "#xxx", "role": "brand.primary", "problem": "fails AA on white", "current_contrast": 3.2, "fix": "#xxx", "fixed_contrast": 4.8}}
+    ],
+    "score": 5
+  }},
   "spacing": {{"analysis": "...", "detected_base": 8, "score": 7, "recommendations": ["..."]}},
   "top_3_priorities": ["...", "...", "..."],
   "confidence": 85
@@ -757,15 +840,34 @@ Analyze these design tokens extracted from a website and compare against industr
 ```"""
 
 
-def build_head_prompt(llm1_analysis: dict, llm2_analysis: dict, rule_calculations: dict) -> str:
-    """Build prompt for HEAD compiler."""
-    return f"""You are a Principal Design Systems Architect compiling analyses from two expert analysts.
+def build_head_prompt(llm1_analysis: dict, llm2_analysis: dict, rule_calculations: dict, semantic_analysis: dict = None) -> str:
+    """Build prompt for HEAD compiler with semantic context."""
+    
+    # Build semantic summary for HEAD
+    semantic_summary = ""
+    if semantic_analysis:
+        brand = semantic_analysis.get("brand", {})
+        brand_primary = brand.get("primary", {}).get("hex", "unknown") if brand.get("primary") else "unknown"
+        brand_secondary = brand.get("secondary", {}).get("hex", "unknown") if brand.get("secondary") else "unknown"
+        
+        text = semantic_analysis.get("text", {})
+        text_primary = text.get("primary", {}).get("hex", "unknown") if text.get("primary") else "unknown"
+        
+        semantic_summary = f"""
+## SEMANTIC COLOR CONTEXT (from Stage 1 Analysis)
+- Brand Primary: {brand_primary}
+- Brand Secondary: {brand_secondary}  
+- Text Primary: {text_primary}
+- Analysis Method: {semantic_analysis.get('summary', {}).get('method', 'rule-based')}
+"""
 
+    return f"""You are a Principal Design Systems Architect compiling analyses from two expert analysts.
+{semantic_summary}
 ## ANALYST 1 FINDINGS:
-{json.dumps(llm1_analysis, indent=2, default=str)[:2000]}
+{json.dumps(llm1_analysis, indent=2, default=str)[:2500]}
 
 ## ANALYST 2 FINDINGS:
-{json.dumps(llm2_analysis, indent=2, default=str)[:2000]}
+{json.dumps(llm2_analysis, indent=2, default=str)[:2500]}
 
 ## RULE-BASED CALCULATIONS:
 - Base font size: {rule_calculations.get('base_font_size', 16)}px
@@ -773,22 +875,30 @@ def build_head_prompt(llm1_analysis: dict, llm2_analysis: dict, rule_calculation
 - Spacing options: 4px grid, 8px grid
 
 ## YOUR TASK:
-1. Compare both analyst perspectives
+1. Compare both analyst perspectives on typography, colors, spacing
 2. Identify agreements and disagreements
-3. Synthesize final recommendations
+3. For COLOR recommendations: Use the semantic context above to make SPECIFIC recommendations per role
+4. Synthesize final recommendations that a designer can act on
 
 ## RESPOND IN JSON FORMAT ONLY:
 ```json
 {{
   "agreements": [{{"topic": "...", "finding": "..."}}],
-  "disagreements": [{{"topic": "...", "resolution": "..."}}],
+  "disagreements": [{{"topic": "...", "analyst1_view": "...", "analyst2_view": "...", "resolution": "..."}}],
   "final_recommendations": {{
     "type_scale": "1.25",
     "type_scale_rationale": "...",
     "spacing_base": "8px",
     "spacing_rationale": "...",
-    "color_improvements": ["..."],
-    "accessibility_fixes": ["..."]
+    "color_recommendations": {{
+      "brand_primary": {{"current": "#xxx", "action": "keep|change", "suggested": "#xxx", "rationale": "..."}},
+      "brand_secondary": {{"current": "#xxx", "action": "keep|change", "suggested": "#xxx", "rationale": "..."}},
+      "text_primary": {{"current": "#xxx", "action": "keep|change", "suggested": "#xxx", "rationale": "..."}},
+      "generate_ramps_for": ["brand.primary", "brand.secondary", "neutral"]
+    }},
+    "accessibility_fixes": [
+      {{"color": "#xxx", "role": "...", "issue": "...", "fix": "#xxx"}}
+    ]
   }},
   "overall_confidence": 85,
   "summary": "..."
@@ -911,8 +1021,9 @@ async def run_stage2_multi_agent(
     mobile_tokens: dict,
     competitors: list[str],
     log_callback: Optional[Callable] = None,
+    semantic_analysis: Optional[dict] = None,
 ) -> dict:
-    """Run the Stage 2 multi-agent analysis."""
+    """Run the Stage 2 multi-agent analysis with semantic context."""
     
     global cost_tracker
     cost_tracker = CostTracker()  # Reset
@@ -923,6 +1034,17 @@ async def run_stage2_multi_agent(
         log_callback("🧠 STAGE 2: MULTI-AGENT ANALYSIS")
         log_callback("=" * 60)
         log_callback("")
+        
+        # Log semantic context
+        if semantic_analysis:
+            brand = semantic_analysis.get("brand", {})
+            brand_primary = brand.get("primary", {}).get("hex", "unknown") if brand.get("primary") else "unknown"
+            log_callback("🧠 SEMANTIC CONTEXT FROM STAGE 1:")
+            log_callback(f"   Brand Primary: {brand_primary}")
+            log_callback(f"   Text Primary: {semantic_analysis.get('text', {}).get('primary', {}).get('hex', 'unknown')}")
+            log_callback(f"   Analysis Method: {semantic_analysis.get('summary', {}).get('method', 'rule-based')}")
+            log_callback("")
+        
         log_callback("📦 LLM CONFIGURATION:")
         
         config = load_agent_config()
@@ -940,11 +1062,12 @@ async def run_stage2_multi_agent(
         log_callback("")
         log_callback("🔄 RUNNING PARALLEL ANALYSIS...")
     
-    # Initial state
+    # Initial state with semantic analysis
     initial_state = {
         "desktop_tokens": desktop_tokens,
         "mobile_tokens": mobile_tokens,
         "competitors": competitors,
+        "semantic_analysis": semantic_analysis,  # NEW: Pass semantic context
         "llm1_analysis": None,
         "llm2_analysis": None,
         "rule_calculations": None,