Spaces:
Sleeping
Sleeping
| """ | |
| Design System Extractor v2 — Main Application | |
| ============================================== | |
| A semi-automated, human-in-the-loop agentic system that reverse-engineers | |
| design systems from live websites. | |
| Usage: | |
| python app.py | |
| """ | |
| import os | |
| import asyncio | |
| import gradio as gr | |
| from datetime import datetime | |
| # Get HF token from environment if available | |
| HF_TOKEN_FROM_ENV = os.getenv("HF_TOKEN", "") | |
| # ============================================================================= | |
| # GLOBAL STATE | |
| # ============================================================================= | |
| current_extraction: dict = {} | |
| user_hf_token: str = "" | |
| # ============================================================================= | |
| # HF TOKEN MANAGEMENT | |
| # ============================================================================= | |
| def set_hf_token(token: str) -> str: | |
| """Set the HF token globally.""" | |
| global user_hf_token | |
| if not token or len(token) < 10: | |
| return "❌ Please enter a valid HuggingFace token" | |
| user_hf_token = token.strip() | |
| os.environ["HF_TOKEN"] = user_hf_token | |
| return "✅ Token saved! You can now use the extractor." | |
| # ============================================================================= | |
| # LAZY IMPORTS (avoid circular imports at startup) | |
| # ============================================================================= | |
| _crawler_module = None | |
| _extractor_module = None | |
| _schema_module = None | |
| def get_crawler(): | |
| global _crawler_module | |
| if _crawler_module is None: | |
| from agents import crawler as _crawler_module | |
| return _crawler_module | |
| def get_extractor(): | |
| global _extractor_module | |
| if _extractor_module is None: | |
| from agents import extractor as _extractor_module | |
| return _extractor_module | |
| def get_schema(): | |
| global _schema_module | |
| if _schema_module is None: | |
| from core import token_schema as _schema_module | |
| return _schema_module | |
| # ============================================================================= | |
| # STAGE 1: URL INPUT & PAGE DISCOVERY | |
| # ============================================================================= | |
| async def discover_site_pages(url: str, progress=gr.Progress()) -> tuple: | |
| """ | |
| Discover pages from a website URL. | |
| Returns tuple of (status_message, pages_dataframe, pages_json) | |
| """ | |
| if not url or not url.startswith(("http://", "https://")): | |
| return "❌ Please enter a valid URL starting with http:// or https://", None, None | |
| progress(0, desc="Initializing browser...") | |
| try: | |
| crawler = get_crawler() | |
| discoverer = crawler.PageDiscoverer() | |
| def update_progress(p): | |
| progress(p, desc=f"Discovering pages... ({int(p*100)}%)") | |
| pages = await discoverer.discover(url, progress_callback=update_progress) | |
| # Format for display | |
| pages_data = [] | |
| for page in pages: | |
| pages_data.append({ | |
| "Select": page.selected, | |
| "URL": page.url, | |
| "Title": page.title or "(No title)", | |
| "Type": page.page_type.value, | |
| "Status": "✓" if not page.error else f"⚠ {page.error}", | |
| }) | |
| # Store for later use | |
| current_extraction["discovered_pages"] = pages | |
| current_extraction["base_url"] = url | |
| status = f"✅ Found {len(pages)} pages. Select the pages you want to extract tokens from." | |
| return status, pages_data, [p.model_dump() for p in pages] | |
| except Exception as e: | |
| import traceback | |
| return f"❌ Error: {str(e)}\n\n{traceback.format_exc()}", None, None | |
| async def start_extraction(pages_selection: list, viewport_choice: str, progress=gr.Progress()) -> tuple: | |
| """ | |
| Start token extraction from selected pages. | |
| Returns tuple of (status, colors_data, typography_data, spacing_data) | |
| """ | |
| if not pages_selection: | |
| return "❌ Please select at least one page", None, None, None | |
| # Get selected URLs | |
| selected_urls = [] | |
| for row in pages_selection: | |
| if row.get("Select", False): | |
| selected_urls.append(row["URL"]) | |
| if not selected_urls: | |
| return "❌ Please select at least one page using the checkboxes", None, None, None | |
| # Determine viewport | |
| schema = get_schema() | |
| viewport = schema.Viewport.DESKTOP if viewport_choice == "Desktop (1440px)" else schema.Viewport.MOBILE | |
| progress(0, desc=f"Starting {viewport.value} extraction...") | |
| try: | |
| extractor_mod = get_extractor() | |
| extractor = extractor_mod.TokenExtractor(viewport=viewport) | |
| def update_progress(p): | |
| progress(p, desc=f"Extracting tokens... ({int(p*100)}%)") | |
| result = await extractor.extract(selected_urls, progress_callback=update_progress) | |
| # Store result | |
| current_extraction[f"{viewport.value}_tokens"] = result | |
| # Format colors for display | |
| colors_data = [] | |
| for color in sorted(result.colors, key=lambda c: -c.frequency)[:50]: | |
| colors_data.append({ | |
| "Accept": True, | |
| "Color": color.value, | |
| "Frequency": color.frequency, | |
| "Context": ", ".join(color.contexts[:3]), | |
| "Contrast (White)": f"{color.contrast_white}:1", | |
| "AA Text": "✓" if color.wcag_aa_small_text else "✗", | |
| "Confidence": color.confidence.value, | |
| }) | |
| # Format typography for display | |
| typography_data = [] | |
| for typo in sorted(result.typography, key=lambda t: -t.frequency)[:30]: | |
| typography_data.append({ | |
| "Accept": True, | |
| "Font": typo.font_family, | |
| "Size": typo.font_size, | |
| "Weight": typo.font_weight, | |
| "Line Height": typo.line_height, | |
| "Elements": ", ".join(typo.elements[:3]), | |
| "Frequency": typo.frequency, | |
| }) | |
| # Format spacing for display | |
| spacing_data = [] | |
| for space in sorted(result.spacing, key=lambda s: s.value_px)[:20]: | |
| spacing_data.append({ | |
| "Accept": True, | |
| "Value": space.value, | |
| "Frequency": space.frequency, | |
| "Context": ", ".join(space.contexts[:2]), | |
| "Fits 8px": "✓" if space.fits_base_8 else "", | |
| "Outlier": "⚠" if space.is_outlier else "", | |
| }) | |
| # Summary | |
| status = f"""✅ Extraction Complete ({viewport.value}) | |
| **Summary:** | |
| - Pages crawled: {len(result.pages_crawled)} | |
| - Colors found: {len(result.colors)} | |
| - Typography styles: {len(result.typography)} | |
| - Spacing values: {len(result.spacing)} | |
| - Font families: {len(result.font_families)} | |
| - Detected spacing base: {result.spacing_base or 'Unknown'}px | |
| - Duration: {result.extraction_duration_ms}ms | |
| """ | |
| if result.warnings: | |
| status += f"\n⚠️ Warnings: {len(result.warnings)}" | |
| if result.errors: | |
| status += f"\n❌ Errors: {len(result.errors)}" | |
| return status, colors_data, typography_data, spacing_data | |
| except Exception as e: | |
| import traceback | |
| return f"❌ Extraction failed: {str(e)}\n\n{traceback.format_exc()}", None, None, None | |
| def export_tokens_json(): | |
| """Export current tokens to JSON.""" | |
| import json | |
| result = {} | |
| if "desktop_tokens" in current_extraction: | |
| desktop = current_extraction["desktop_tokens"] | |
| result["desktop"] = { | |
| "colors": [c.model_dump() for c in desktop.colors], | |
| "typography": [t.model_dump() for t in desktop.typography], | |
| "spacing": [s.model_dump() for s in desktop.spacing], | |
| "metadata": desktop.summary(), | |
| } | |
| if "mobile_tokens" in current_extraction: | |
| mobile = current_extraction["mobile_tokens"] | |
| result["mobile"] = { | |
| "colors": [c.model_dump() for c in mobile.colors], | |
| "typography": [t.model_dump() for t in mobile.typography], | |
| "spacing": [s.model_dump() for s in mobile.spacing], | |
| "metadata": mobile.summary(), | |
| } | |
| if not result: | |
| return '{"error": "No tokens extracted yet. Please run extraction first."}' | |
| return json.dumps(result, indent=2, default=str) | |
| # ============================================================================= | |
| # UI BUILDING | |
| # ============================================================================= | |
| def create_ui(): | |
| """Create the Gradio interface.""" | |
| with gr.Blocks( | |
| title="Design System Extractor v2", | |
| theme=gr.themes.Soft(), | |
| ) as app: | |
| # Header | |
| gr.Markdown(""" | |
| # 🎨 Design System Extractor v2 | |
| **Reverse-engineer design systems from live websites.** | |
| Extract colors, typography, and spacing tokens from any website and export to Figma-compatible JSON. | |
| --- | |
| """) | |
| # ================================================================= | |
| # CONFIGURATION SECTION | |
| # ================================================================= | |
| with gr.Accordion("⚙️ Configuration", open=not bool(HF_TOKEN_FROM_ENV)): | |
| gr.Markdown(""" | |
| **HuggingFace Token** is required for AI-powered features (Agent 2-4). | |
| Get your token at: [huggingface.co/settings/tokens](https://huggingface.co/settings/tokens) | |
| *Note: Basic extraction (Agent 1) works without a token.* | |
| """) | |
| with gr.Row(): | |
| hf_token_input = gr.Textbox( | |
| label="HuggingFace Token", | |
| placeholder="hf_xxxxxxxxxxxxxxxxxxxx", | |
| type="password", | |
| scale=4, | |
| value=HF_TOKEN_FROM_ENV if HF_TOKEN_FROM_ENV else "", | |
| ) | |
| save_token_btn = gr.Button("💾 Save Token", scale=1) | |
| token_status = gr.Markdown( | |
| "✅ Token loaded from environment" if HF_TOKEN_FROM_ENV else "⏳ Enter your HF token to enable all features" | |
| ) | |
| save_token_btn.click( | |
| fn=set_hf_token, | |
| inputs=[hf_token_input], | |
| outputs=[token_status], | |
| ) | |
| # ================================================================= | |
| # STAGE 1: URL Input & Discovery | |
| # ================================================================= | |
| with gr.Accordion("📍 Stage 1: Website Discovery", open=True): | |
| gr.Markdown(""" | |
| **Step 1:** Enter your website URL and discover pages. | |
| The system will automatically find and classify pages for extraction. | |
| """) | |
| with gr.Row(): | |
| url_input = gr.Textbox( | |
| label="Website URL", | |
| placeholder="https://example.com", | |
| scale=4, | |
| ) | |
| discover_btn = gr.Button("🔍 Discover Pages", variant="primary", scale=1) | |
| discovery_status = gr.Markdown("") | |
| pages_table = gr.Dataframe( | |
| headers=["Select", "URL", "Title", "Type", "Status"], | |
| datatype=["bool", "str", "str", "str", "str"], | |
| interactive=True, | |
| label="Discovered Pages", | |
| visible=False, | |
| ) | |
| pages_json = gr.JSON(visible=False) | |
| # ================================================================= | |
| # STAGE 2: Extraction | |
| # ================================================================= | |
| with gr.Accordion("🔬 Stage 2: Token Extraction", open=False): | |
| gr.Markdown(""" | |
| **Step 2:** Select pages and viewport, then extract design tokens. | |
| """) | |
| with gr.Row(): | |
| viewport_radio = gr.Radio( | |
| choices=["Desktop (1440px)", "Mobile (375px)"], | |
| value="Desktop (1440px)", | |
| label="Viewport", | |
| ) | |
| extract_btn = gr.Button("🚀 Extract Tokens", variant="primary") | |
| extraction_status = gr.Markdown("") | |
| with gr.Tabs(): | |
| with gr.Tab("🎨 Colors"): | |
| colors_table = gr.Dataframe( | |
| headers=["Accept", "Color", "Frequency", "Context", "Contrast (White)", "AA Text", "Confidence"], | |
| datatype=["bool", "str", "number", "str", "str", "str", "str"], | |
| interactive=True, | |
| label="Extracted Colors", | |
| ) | |
| with gr.Tab("📝 Typography"): | |
| typography_table = gr.Dataframe( | |
| headers=["Accept", "Font", "Size", "Weight", "Line Height", "Elements", "Frequency"], | |
| datatype=["bool", "str", "str", "number", "str", "str", "number"], | |
| interactive=True, | |
| label="Extracted Typography", | |
| ) | |
| with gr.Tab("📏 Spacing"): | |
| spacing_table = gr.Dataframe( | |
| headers=["Accept", "Value", "Frequency", "Context", "Fits 8px", "Outlier"], | |
| datatype=["bool", "str", "number", "str", "str", "str"], | |
| interactive=True, | |
| label="Extracted Spacing", | |
| ) | |
| # ================================================================= | |
| # STAGE 3: Export | |
| # ================================================================= | |
| with gr.Accordion("📦 Stage 3: Export", open=False): | |
| gr.Markdown(""" | |
| **Step 3:** Review and export your design tokens. | |
| """) | |
| with gr.Row(): | |
| export_btn = gr.Button("📥 Export JSON", variant="secondary") | |
| export_output = gr.Code( | |
| label="Exported Tokens (JSON)", | |
| language="json", | |
| lines=20, | |
| ) | |
| # ================================================================= | |
| # EVENT HANDLERS | |
| # ================================================================= | |
| # Discovery | |
| discover_btn.click( | |
| fn=discover_site_pages, | |
| inputs=[url_input], | |
| outputs=[discovery_status, pages_table, pages_json], | |
| ).then( | |
| fn=lambda: gr.update(visible=True), | |
| outputs=[pages_table], | |
| ) | |
| # Extraction | |
| extract_btn.click( | |
| fn=start_extraction, | |
| inputs=[pages_table, viewport_radio], | |
| outputs=[extraction_status, colors_table, typography_table, spacing_table], | |
| ) | |
| # Export | |
| export_btn.click( | |
| fn=export_tokens_json, | |
| outputs=[export_output], | |
| ) | |
| # ================================================================= | |
| # FOOTER | |
| # ================================================================= | |
| gr.Markdown(""" | |
| --- | |
| **Design System Extractor v2** | Built with LangGraph + Gradio + HuggingFace | |
| *A semi-automated co-pilot for design system recovery and modernization.* | |
| **Models:** Microsoft Phi (Normalizer) • Meta Llama (Advisor) • Mistral Codestral (Generator) | |
| """) | |
| return app | |
| # ============================================================================= | |
| # MAIN | |
| # ============================================================================= | |
| if __name__ == "__main__": | |
| app = create_ui() | |
| app.launch( | |
| server_name="0.0.0.0", | |
| server_port=7860, | |
| ) | |