Upload agent_2_website_inspector.py

agents/agent_2_website_inspector.py

@@ -0,0 +1,181 @@
+"""
+Agent 2: Website Inspector
+Captures screenshots AND extracts DOM elements from the live website.
+Gets computed styles for each element for comparison with Figma.
+"""
+
+from typing import Dict, Any
+import sys
+import os
+import asyncio
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+from utils.dom_element_extractor import DOMElementExtractor
+
+
+async def _capture_and_extract_async(
+    website_url: str,
+    output_dir: str,
+    execution_id: str,
+    desktop_width: int = 1440,
+    mobile_width: int = 375
+) -> tuple:
+    """Async function to capture screenshots and extract elements."""
+    from playwright.async_api import async_playwright
+    from pathlib import Path
+    
+    Path(output_dir).mkdir(parents=True, exist_ok=True)
+    
+    screenshots = {}
+    dimensions = {}
+    all_elements = {}
+    all_summaries = {}
+    
+    async with async_playwright() as p:
+        browser = await p.chromium.launch(headless=True)
+        
+        try:
+            # Desktop capture and extraction
+            print(f"\n  📱 Processing desktop ({desktop_width}px)...")
+            page = await browser.new_page(viewport={"width": desktop_width, "height": 1080})
+            await page.goto(website_url, wait_until="networkidle", timeout=60000)
+            
+            # Wait longer for CSS/JS to fully load and apply styles
+            await page.wait_for_timeout(3000)
+            
+            # Scroll to trigger any lazy-loaded styles
+            await page.evaluate("window.scrollTo(0, document.body.scrollHeight)")
+            await page.wait_for_timeout(500)
+            await page.evaluate("window.scrollTo(0, 0)")
+            await page.wait_for_timeout(500)
+            
+            # Get full page height
+            desktop_height = await page.evaluate("() => document.documentElement.scrollHeight")
+            await page.set_viewport_size({"width": desktop_width, "height": desktop_height})
+            
+            # Screenshot
+            desktop_path = f"{output_dir}/desktop_{execution_id}.png"
+            await page.screenshot(path=desktop_path, full_page=True)
+            screenshots["desktop"] = desktop_path
+            dimensions["desktop"] = {"width": desktop_width, "height": desktop_height}
+            print(f"     📸 Screenshot: {desktop_width}x{desktop_height}px")
+            
+            # Extract elements
+            extractor = DOMElementExtractor()
+            elements = await extractor.extract_from_page_async(page, "desktop")
+            all_elements["desktop"] = [e.to_dict() for e in elements]
+            all_summaries["desktop"] = extractor.summarize()
+            print(f"     🔍 Elements: {len(elements)} extracted")
+            
+            await page.close()
+            
+            # Mobile capture and extraction
+            print(f"\n  📱 Processing mobile ({mobile_width}px)...")
+            page = await browser.new_page(viewport={"width": mobile_width, "height": 812})
+            await page.set_extra_http_headers({
+                "User-Agent": "Mozilla/5.0 (iPhone; CPU iPhone OS 15_0 like Mac OS X) AppleWebKit/605.1.15"
+            })
+            await page.goto(website_url, wait_until="networkidle", timeout=60000)
+            await page.wait_for_timeout(2000)
+            
+            mobile_height = await page.evaluate("() => document.documentElement.scrollHeight")
+            await page.set_viewport_size({"width": mobile_width, "height": mobile_height})
+            
+            # Screenshot
+            mobile_path = f"{output_dir}/mobile_{execution_id}.png"
+            await page.screenshot(path=mobile_path, full_page=True)
+            screenshots["mobile"] = mobile_path
+            dimensions["mobile"] = {"width": mobile_width, "height": mobile_height}
+            print(f"     📸 Screenshot: {mobile_width}x{mobile_height}px")
+            
+            # Extract elements
+            extractor = DOMElementExtractor()
+            elements = await extractor.extract_from_page_async(page, "mobile")
+            all_elements["mobile"] = [e.to_dict() for e in elements]
+            all_summaries["mobile"] = extractor.summarize()
+            print(f"     🔍 Elements: {len(elements)} extracted")
+            
+            await page.close()
+            
+        finally:
+            await browser.close()
+    
+    return screenshots, dimensions, all_elements, all_summaries
+
+
+def agent_2_node(state: Dict[str, Any]) -> Dict[str, Any]:
+    """
+    Capture screenshots AND extract DOM elements from website.
+    
+    This agent:
+    1. Uses Playwright to load the website
+    2. Captures full-page screenshots at Desktop and Mobile viewports
+    3. Extracts ALL DOM elements with computed CSS styles
+    4. Stores paths, dimensions, and elements in state
+    """
+    print("\n" + "="*60)
+    print("🌐 Agent 2: Website Inspector - Website Analysis")
+    print("="*60)
+    
+    website_url = state.get("website_url", "")
+    execution_id = state.get("execution_id", "")
+    figma_dims = state.get("figma_dimensions", {})
+    logs = state.get("logs", [])
+    
+    try:
+        # Get desktop width from Figma (if available) for consistent comparison
+        desktop_width = 1440
+        if "desktop" in figma_dims:
+            desktop_width = figma_dims["desktop"].get("width", 1440)
+        
+        print(f"  🔗 URL: {website_url}")
+        print(f"  📐 Desktop width: {desktop_width}px (matching Figma)")
+        
+        # Capture and extract
+        screenshots, dimensions, all_elements, all_summaries = asyncio.run(
+            _capture_and_extract_async(
+                website_url=website_url,
+                output_dir="data/website",
+                execution_id=execution_id,
+                desktop_width=desktop_width,
+                mobile_width=375
+            )
+        )
+        
+        if not screenshots:
+            raise ValueError("Failed to capture any screenshots")
+        
+        # Log results
+        print(f"\n  ✅ Website analysis complete")
+        for viewport in screenshots:
+            dims = dimensions.get(viewport, {})
+            summary = all_summaries.get(viewport, {})
+            logs.append(f"📸 Website {viewport}: {dims.get('width', '?')}x{dims.get('height', '?')}px")
+            logs.append(f"🔍 Website {viewport}: {summary.get('total_elements', 0)} elements")
+        
+        total_elements = sum(s.get('total_elements', 0) for s in all_summaries.values())
+        print(f"     Total: {len(screenshots)} screenshots, {total_elements} elements")
+        
+        return {
+            "website_screenshots": screenshots,
+            "website_dimensions": dimensions,
+            "website_elements": all_elements,
+            "website_element_summary": all_summaries,
+            "status": "website_captured",
+            "logs": logs
+        }
+        
+    except Exception as e:
+        error_msg = f"Failed to analyze website: {str(e)}"
+        print(f"\n  ❌ {error_msg}")
+        logs.append(f"❌ {error_msg}")
+        
+        import traceback
+        traceback.print_exc()
+        
+        return {
+            "status": "website_capture_failed",
+            "error_message": error_msg,
+            "logs": logs
+        }