riazmo commited on
Commit
7e6bfe9
Β·
verified Β·
1 Parent(s): 77580de

Upload app.py

Browse files
Files changed (1) hide show
  1. app.py +1044 -0
app.py ADDED
@@ -0,0 +1,1044 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Design System Extractor v2 β€” Main Application
3
+ ==============================================
4
+
5
+ Flow:
6
+ 1. User enters URL
7
+ 2. Agent 1 discovers pages β†’ User confirms
8
+ 3. Agent 1 extracts tokens (Desktop + Mobile)
9
+ 4. Agent 2 normalizes tokens
10
+ 5. Stage 1 UI: User reviews tokens (accept/reject, Desktop↔Mobile toggle)
11
+ 6. Agent 3 proposes upgrades
12
+ 7. Stage 2 UI: User selects options with live preview
13
+ 8. Agent 4 generates JSON
14
+ 9. Stage 3 UI: User exports
15
+ """
16
+
17
+ import os
18
+ import asyncio
19
+ import json
20
+ import gradio as gr
21
+ from datetime import datetime
22
+ from typing import Optional
23
+
24
+ # Get HF token from environment
25
+ HF_TOKEN_FROM_ENV = os.getenv("HF_TOKEN", "")
26
+
27
+ # =============================================================================
28
+ # GLOBAL STATE
29
+ # =============================================================================
30
+
31
+ class AppState:
32
+ """Global application state."""
33
+ def __init__(self):
34
+ self.reset()
35
+
36
+ def reset(self):
37
+ self.discovered_pages = []
38
+ self.base_url = ""
39
+ self.desktop_raw = None # ExtractedTokens
40
+ self.mobile_raw = None # ExtractedTokens
41
+ self.desktop_normalized = None # NormalizedTokens
42
+ self.mobile_normalized = None # NormalizedTokens
43
+ self.upgrade_recommendations = None # UpgradeRecommendations
44
+ self.selected_upgrades = {} # User selections
45
+ self.logs = []
46
+
47
+ def log(self, message: str):
48
+ timestamp = datetime.now().strftime("%H:%M:%S")
49
+ self.logs.append(f"[{timestamp}] {message}")
50
+ if len(self.logs) > 100:
51
+ self.logs.pop(0)
52
+
53
+ def get_logs(self) -> str:
54
+ return "\n".join(self.logs)
55
+
56
+ state = AppState()
57
+
58
+
59
+ # =============================================================================
60
+ # LAZY IMPORTS
61
+ # =============================================================================
62
+
63
+ def get_crawler():
64
+ import agents.crawler
65
+ return agents.crawler
66
+
67
+ def get_extractor():
68
+ import agents.extractor
69
+ return agents.extractor
70
+
71
+ def get_normalizer():
72
+ import agents.normalizer
73
+ return agents.normalizer
74
+
75
+ def get_advisor():
76
+ import agents.advisor
77
+ return agents.advisor
78
+
79
+ def get_schema():
80
+ import core.token_schema
81
+ return core.token_schema
82
+
83
+
84
+ # =============================================================================
85
+ # PHASE 1: DISCOVER PAGES
86
+ # =============================================================================
87
+
88
+ async def discover_pages(url: str, progress=gr.Progress()):
89
+ """Discover pages from URL."""
90
+ state.reset()
91
+
92
+ if not url or not url.startswith(("http://", "https://")):
93
+ return "❌ Please enter a valid URL", "", None
94
+
95
+ state.log(f"πŸš€ Starting discovery for: {url}")
96
+ progress(0.1, desc="πŸ” Discovering pages...")
97
+
98
+ try:
99
+ crawler = get_crawler()
100
+ discoverer = crawler.PageDiscoverer()
101
+
102
+ pages = await discoverer.discover(url)
103
+
104
+ state.discovered_pages = pages
105
+ state.base_url = url
106
+
107
+ state.log(f"βœ… Found {len(pages)} pages")
108
+
109
+ # Format for display
110
+ pages_data = []
111
+ for page in pages:
112
+ pages_data.append([
113
+ True, # Selected by default
114
+ page.url,
115
+ page.title if page.title else "(No title)",
116
+ page.page_type.value,
117
+ "βœ“" if not page.error else f"⚠ {page.error}"
118
+ ])
119
+
120
+ progress(1.0, desc="βœ… Discovery complete!")
121
+
122
+ status = f"βœ… Found {len(pages)} pages. Review and click 'Extract Tokens' to continue."
123
+
124
+ return status, state.get_logs(), pages_data
125
+
126
+ except Exception as e:
127
+ import traceback
128
+ state.log(f"❌ Error: {str(e)}")
129
+ return f"❌ Error: {str(e)}", state.get_logs(), None
130
+
131
+
132
+ # =============================================================================
133
+ # PHASE 2: EXTRACT TOKENS
134
+ # =============================================================================
135
+
136
+ async def extract_tokens(pages_data, progress=gr.Progress()):
137
+ """Extract tokens from selected pages (both viewports)."""
138
+
139
+ state.log(f"πŸ“₯ Received pages_data type: {type(pages_data)}")
140
+
141
+ if pages_data is None:
142
+ return "❌ Please discover pages first", state.get_logs(), None, None
143
+
144
+ # Get selected URLs - handle pandas DataFrame
145
+ selected_urls = []
146
+
147
+ try:
148
+ # Check if it's a pandas DataFrame
149
+ if hasattr(pages_data, 'iterrows'):
150
+ state.log(f"πŸ“₯ DataFrame with {len(pages_data)} rows, columns: {list(pages_data.columns)}")
151
+
152
+ for idx, row in pages_data.iterrows():
153
+ # Get values by column name or position
154
+ try:
155
+ # Try column names first
156
+ is_selected = row.get('Select', row.iloc[0] if len(row) > 0 else False)
157
+ url = row.get('URL', row.iloc[1] if len(row) > 1 else '')
158
+ except:
159
+ # Fallback to positional
160
+ is_selected = row.iloc[0] if len(row) > 0 else False
161
+ url = row.iloc[1] if len(row) > 1 else ''
162
+
163
+ if is_selected and url:
164
+ selected_urls.append(url)
165
+
166
+ # If it's a dict (Gradio sometimes sends this)
167
+ elif isinstance(pages_data, dict):
168
+ state.log(f"πŸ“₯ Dict with keys: {list(pages_data.keys())}")
169
+ data = pages_data.get('data', [])
170
+ for row in data:
171
+ if isinstance(row, (list, tuple)) and len(row) >= 2 and row[0]:
172
+ selected_urls.append(row[1])
173
+
174
+ # If it's a list
175
+ elif isinstance(pages_data, (list, tuple)):
176
+ state.log(f"πŸ“₯ List with {len(pages_data)} items")
177
+ for row in pages_data:
178
+ if isinstance(row, (list, tuple)) and len(row) >= 2 and row[0]:
179
+ selected_urls.append(row[1])
180
+
181
+ except Exception as e:
182
+ state.log(f"❌ Error parsing pages_data: {str(e)}")
183
+ import traceback
184
+ state.log(traceback.format_exc())
185
+
186
+ state.log(f"πŸ“‹ Found {len(selected_urls)} selected URLs")
187
+
188
+ # If still no URLs, try using stored discovered pages
189
+ if not selected_urls and state.discovered_pages:
190
+ state.log("⚠️ No URLs from table, using all discovered pages")
191
+ selected_urls = [p.url for p in state.discovered_pages if not p.error][:10]
192
+
193
+ if not selected_urls:
194
+ return "❌ No pages selected. Please select pages or rediscover.", state.get_logs(), None, None
195
+
196
+ # Limit to 10 pages for performance
197
+ selected_urls = selected_urls[:10]
198
+
199
+ state.log(f"πŸ“‹ Extracting from {len(selected_urls)} pages:")
200
+ for url in selected_urls[:3]:
201
+ state.log(f" β€’ {url}")
202
+ if len(selected_urls) > 3:
203
+ state.log(f" ... and {len(selected_urls) - 3} more")
204
+
205
+ progress(0.05, desc="πŸš€ Starting extraction...")
206
+
207
+ try:
208
+ schema = get_schema()
209
+ extractor_mod = get_extractor()
210
+ normalizer_mod = get_normalizer()
211
+
212
+ # === DESKTOP EXTRACTION ===
213
+ state.log("")
214
+ state.log("πŸ–₯️ DESKTOP EXTRACTION (1440px)")
215
+ progress(0.1, desc="πŸ–₯️ Extracting desktop tokens...")
216
+
217
+ desktop_extractor = extractor_mod.TokenExtractor(viewport=schema.Viewport.DESKTOP)
218
+
219
+ def desktop_progress(p):
220
+ progress(0.1 + (p * 0.35), desc=f"πŸ–₯️ Desktop... {int(p*100)}%")
221
+
222
+ state.desktop_raw = await desktop_extractor.extract(selected_urls, progress_callback=desktop_progress)
223
+
224
+ state.log(f" Raw: {len(state.desktop_raw.colors)} colors, {len(state.desktop_raw.typography)} typography, {len(state.desktop_raw.spacing)} spacing")
225
+
226
+ # Normalize desktop
227
+ state.log(" Normalizing...")
228
+ state.desktop_normalized = normalizer_mod.normalize_tokens(state.desktop_raw)
229
+ state.log(f" Normalized: {len(state.desktop_normalized.colors)} colors, {len(state.desktop_normalized.typography)} typography, {len(state.desktop_normalized.spacing)} spacing")
230
+
231
+ # === MOBILE EXTRACTION ===
232
+ state.log("")
233
+ state.log("πŸ“± MOBILE EXTRACTION (375px)")
234
+ progress(0.5, desc="πŸ“± Extracting mobile tokens...")
235
+
236
+ mobile_extractor = extractor_mod.TokenExtractor(viewport=schema.Viewport.MOBILE)
237
+
238
+ def mobile_progress(p):
239
+ progress(0.5 + (p * 0.35), desc=f"πŸ“± Mobile... {int(p*100)}%")
240
+
241
+ state.mobile_raw = await mobile_extractor.extract(selected_urls, progress_callback=mobile_progress)
242
+
243
+ state.log(f" Raw: {len(state.mobile_raw.colors)} colors, {len(state.mobile_raw.typography)} typography, {len(state.mobile_raw.spacing)} spacing")
244
+
245
+ # Normalize mobile
246
+ state.log(" Normalizing...")
247
+ state.mobile_normalized = normalizer_mod.normalize_tokens(state.mobile_raw)
248
+ state.log(f" Normalized: {len(state.mobile_normalized.colors)} colors, {len(state.mobile_normalized.typography)} typography, {len(state.mobile_normalized.spacing)} spacing")
249
+
250
+ progress(0.95, desc="πŸ“Š Preparing results...")
251
+
252
+ # Format results for Stage 1 UI
253
+ desktop_data = format_tokens_for_display(state.desktop_normalized)
254
+ mobile_data = format_tokens_for_display(state.mobile_normalized)
255
+
256
+ state.log("")
257
+ state.log("=" * 50)
258
+ state.log("βœ… EXTRACTION COMPLETE!")
259
+ state.log("=" * 50)
260
+
261
+ progress(1.0, desc="βœ… Complete!")
262
+
263
+ status = f"""## βœ… Extraction Complete!
264
+
265
+ | Viewport | Colors | Typography | Spacing |
266
+ |----------|--------|------------|---------|
267
+ | Desktop | {len(state.desktop_normalized.colors)} | {len(state.desktop_normalized.typography)} | {len(state.desktop_normalized.spacing)} |
268
+ | Mobile | {len(state.mobile_normalized.colors)} | {len(state.mobile_normalized.typography)} | {len(state.mobile_normalized.spacing)} |
269
+
270
+ **Next:** Review the tokens below. Accept or reject, then proceed to Stage 2.
271
+ """
272
+
273
+ return status, state.get_logs(), desktop_data, mobile_data
274
+
275
+ except Exception as e:
276
+ import traceback
277
+ state.log(f"❌ Error: {str(e)}")
278
+ state.log(traceback.format_exc())
279
+ return f"❌ Error: {str(e)}", state.get_logs(), None, None
280
+
281
+
282
+ def format_tokens_for_display(normalized) -> dict:
283
+ """Format normalized tokens for Gradio display."""
284
+ if normalized is None:
285
+ return {"colors": [], "typography": [], "spacing": []}
286
+
287
+ # Colors are now a dict
288
+ colors = []
289
+ color_items = list(normalized.colors.values()) if isinstance(normalized.colors, dict) else normalized.colors
290
+ for c in sorted(color_items, key=lambda x: -x.frequency)[:50]:
291
+ colors.append([
292
+ True, # Accept checkbox
293
+ c.value,
294
+ c.suggested_name or "",
295
+ c.frequency,
296
+ c.confidence.value if c.confidence else "medium",
297
+ f"{c.contrast_white:.1f}:1" if c.contrast_white else "N/A",
298
+ "βœ“" if c.wcag_aa_small_text else "βœ—",
299
+ ", ".join(c.contexts[:2]) if c.contexts else "",
300
+ ])
301
+
302
+ # Typography
303
+ typography = []
304
+ typo_items = list(normalized.typography.values()) if isinstance(normalized.typography, dict) else normalized.typography
305
+ for t in sorted(typo_items, key=lambda x: -x.frequency)[:30]:
306
+ typography.append([
307
+ True, # Accept checkbox
308
+ t.font_family,
309
+ t.font_size,
310
+ str(t.font_weight),
311
+ t.line_height or "",
312
+ t.suggested_name or "",
313
+ t.frequency,
314
+ t.confidence.value if t.confidence else "medium",
315
+ ])
316
+
317
+ # Spacing
318
+ spacing = []
319
+ spacing_items = list(normalized.spacing.values()) if isinstance(normalized.spacing, dict) else normalized.spacing
320
+ for s in sorted(spacing_items, key=lambda x: x.value_px)[:20]:
321
+ spacing.append([
322
+ True, # Accept checkbox
323
+ s.value,
324
+ f"{s.value_px}px",
325
+ s.suggested_name or "",
326
+ s.frequency,
327
+ "βœ“" if s.fits_base_8 else "",
328
+ s.confidence.value if s.confidence else "medium",
329
+ ])
330
+
331
+ return {
332
+ "colors": colors,
333
+ "typography": typography,
334
+ "spacing": spacing,
335
+ }
336
+
337
+
338
+ def switch_viewport(viewport: str):
339
+ """Switch between desktop and mobile view."""
340
+ if viewport == "Desktop (1440px)":
341
+ data = format_tokens_for_display(state.desktop_normalized)
342
+ else:
343
+ data = format_tokens_for_display(state.mobile_normalized)
344
+
345
+ return data["colors"], data["typography"], data["spacing"]
346
+
347
+
348
+ # =============================================================================
349
+ # STAGE 2: AI ANALYSIS
350
+ # =============================================================================
351
+
352
+ async def run_stage2_analysis(progress=gr.Progress()):
353
+ """Run Agent 3 analysis on extracted tokens."""
354
+
355
+ if not state.desktop_normalized or not state.mobile_normalized:
356
+ return ("❌ Please complete Stage 1 first", "", None, None, "", "", "")
357
+
358
+ state.log("")
359
+ state.log("=" * 50)
360
+ state.log("🧠 STAGE 2: AI-POWERED ANALYSIS")
361
+ state.log("=" * 50)
362
+
363
+ progress(0.1, desc="πŸ€– Starting AI analysis...")
364
+
365
+ try:
366
+ advisor_mod = get_advisor()
367
+
368
+ # Run analysis with logging
369
+ state.log("πŸ” Analyzing design patterns...")
370
+ progress(0.3, desc="πŸ” Analyzing patterns...")
371
+
372
+ recommendations = await advisor_mod.analyze_design_system(
373
+ desktop_tokens=state.desktop_normalized,
374
+ mobile_tokens=state.mobile_normalized,
375
+ log_callback=state.log,
376
+ )
377
+
378
+ state.upgrade_recommendations = recommendations
379
+
380
+ progress(0.9, desc="πŸ“Š Preparing recommendations...")
381
+
382
+ # Format typography comparison table
383
+ typography_data = format_typography_comparison(recommendations)
384
+
385
+ # Format spacing comparison table
386
+ spacing_data = format_spacing_comparison(recommendations)
387
+
388
+ # Format color ramps display
389
+ color_ramps_md = format_color_ramps_display(recommendations)
390
+
391
+ # Format radius display
392
+ radius_md = format_radius_display()
393
+
394
+ # Format shadows display
395
+ shadows_md = format_shadows_display()
396
+
397
+ state.log("βœ… Analysis complete!")
398
+ progress(1.0, desc="βœ… Complete!")
399
+
400
+ status = f"""## 🧠 AI Analysis Complete!
401
+
402
+ ### LLM Recommendation
403
+ {recommendations.llm_rationale if recommendations.llm_rationale else "Using rule-based recommendations."}
404
+
405
+ ### Detected Patterns
406
+ {chr(10).join(['β€’ ' + p for p in recommendations.detected_patterns]) if recommendations.detected_patterns else 'β€’ Standard design patterns detected'}
407
+
408
+ **Review the options below and select your preferences.**
409
+ """
410
+
411
+ return (status, state.get_logs(), typography_data, spacing_data,
412
+ color_ramps_md, radius_md, shadows_md)
413
+
414
+ except Exception as e:
415
+ import traceback
416
+ state.log(f"❌ Error: {str(e)}")
417
+ state.log(traceback.format_exc())
418
+ return (f"❌ Analysis failed: {str(e)}", state.get_logs(), None, None, "", "", "")
419
+
420
+
421
+ def format_typography_comparison(recommendations) -> list:
422
+ """Format typography comparison table."""
423
+ if not state.desktop_normalized:
424
+ return []
425
+
426
+ # Get current typography sorted by size
427
+ current_typo = list(state.desktop_normalized.typography.values())
428
+ current_typo.sort(key=lambda t: -float(t.font_size.replace('px', '').replace('rem', '').replace('em', '') or 16))
429
+
430
+ # Get base font size (most common around 14-18px)
431
+ sizes = [float(t.font_size.replace('px', '').replace('rem', '').replace('em', '') or 16) for t in current_typo]
432
+ base = 16
433
+ for s in sizes:
434
+ if 14 <= s <= 18:
435
+ base = s
436
+ break
437
+
438
+ # Generate scales
439
+ scale_12 = generate_type_scale_values(base, 1.2)
440
+ scale_125 = generate_type_scale_values(base, 1.25)
441
+ scale_133 = generate_type_scale_values(base, 1.333)
442
+
443
+ # Build comparison table
444
+ elements = ["Display", "H1", "H2", "H3", "Body", "Caption", "Small"]
445
+ data = []
446
+
447
+ for i, elem in enumerate(elements):
448
+ current = f"{int(sizes[i])}px" if i < len(sizes) else "β€”"
449
+ s12 = f"{scale_12[i]}px" if i < len(scale_12) else "β€”"
450
+ s125 = f"{scale_125[i]}px" if i < len(scale_125) else "β€”"
451
+ s133 = f"{scale_133[i]}px" if i < len(scale_133) else "β€”"
452
+ data.append([elem, current, s12, s125, s133])
453
+
454
+ return data
455
+
456
+
457
+ def generate_type_scale_values(base: float, ratio: float) -> list:
458
+ """Generate type scale values."""
459
+ # Going up from base
460
+ scales = []
461
+ for i in range(4, -1, -1): # 4 sizes up
462
+ scales.append(int(base * (ratio ** i)))
463
+ # Base
464
+ scales.append(int(base))
465
+ # Going down from base
466
+ for i in range(1, 3): # 2 sizes down
467
+ scales.append(int(base / (ratio ** i)))
468
+ return scales
469
+
470
+
471
+ def format_spacing_comparison(recommendations) -> list:
472
+ """Format spacing comparison table."""
473
+ if not state.desktop_normalized:
474
+ return []
475
+
476
+ # Get current spacing
477
+ current_spacing = list(state.desktop_normalized.spacing.values())
478
+ current_spacing.sort(key=lambda s: s.value_px)
479
+
480
+ data = []
481
+ for s in current_spacing[:10]:
482
+ current = f"{s.value_px}px"
483
+ grid_8 = f"{snap_to_grid(s.value_px, 8)}px"
484
+ grid_4 = f"{snap_to_grid(s.value_px, 4)}px"
485
+
486
+ # Mark if value fits
487
+ if s.value_px == snap_to_grid(s.value_px, 8):
488
+ grid_8 += " βœ“"
489
+ if s.value_px == snap_to_grid(s.value_px, 4):
490
+ grid_4 += " βœ“"
491
+
492
+ data.append([current, grid_8, grid_4])
493
+
494
+ return data
495
+
496
+
497
+ def snap_to_grid(value: float, base: int) -> int:
498
+ """Snap value to grid."""
499
+ return round(value / base) * base
500
+
501
+
502
+ def format_color_ramps_display(recommendations) -> str:
503
+ """Format color ramps for display."""
504
+ if not recommendations or not recommendations.color_ramps:
505
+ return "No color ramps to generate."
506
+
507
+ lines = []
508
+ for opt in recommendations.color_ramps:
509
+ base = opt.values.get("base_color", "")
510
+ role = opt.values.get("role", "unknown")
511
+ ramp = opt.values.get("ramp", {})
512
+
513
+ # Get a few shades for preview
514
+ shades = []
515
+ for shade in ["50", "200", "500", "700", "900"]:
516
+ color = ramp.get(f"{role}.{shade}", "?")
517
+ shades.append(f"`{color}`")
518
+
519
+ lines.append(f"**{role.title()}** (base: {base})")
520
+ lines.append(f" {' β†’ '.join(shades)}")
521
+ lines.append("")
522
+
523
+ return "\n".join(lines) if lines else "No color ramps to generate."
524
+
525
+
526
+ def format_radius_display() -> str:
527
+ """Format radius display."""
528
+ if not state.desktop_normalized or not state.desktop_normalized.radius:
529
+ return "*No border radius values detected.*"
530
+
531
+ radii = list(state.desktop_normalized.radius.values())
532
+ values = [r.value for r in radii[:5]]
533
+
534
+ return f"**Detected:** {', '.join(values)}\n\n*Radius values will be preserved as-is.*"
535
+
536
+
537
+ def format_shadows_display() -> str:
538
+ """Format shadows display."""
539
+ if not state.desktop_normalized or not state.desktop_normalized.shadows:
540
+ return "*No shadow values detected.*"
541
+
542
+ count = len(state.desktop_normalized.shadows)
543
+ return f"**Detected:** {count} shadow style(s)\n\n*Shadow values will be preserved as-is.*"
544
+
545
+
546
+ def apply_selected_upgrades(type_choice: str, spacing_choice: str, apply_ramps: bool):
547
+ """Apply selected upgrade options."""
548
+ if not state.upgrade_recommendations:
549
+ return "❌ Run analysis first", ""
550
+
551
+ state.log("✨ Applying selected upgrades...")
552
+
553
+ # Store selections
554
+ state.selected_upgrades = {
555
+ "type_scale": type_choice,
556
+ "spacing": spacing_choice,
557
+ "color_ramps": apply_ramps,
558
+ }
559
+
560
+ state.log(f" Type Scale: {type_choice}")
561
+ state.log(f" Spacing: {spacing_choice}")
562
+ state.log(f" Color Ramps: {'Yes' if apply_ramps else 'No'}")
563
+
564
+ state.log("βœ… Upgrades applied! Proceed to Stage 3 for export.")
565
+
566
+ return "βœ… Upgrades applied! Proceed to Stage 3 to export.", state.get_logs()
567
+
568
+
569
+ def export_stage1_json():
570
+ """Export Stage 1 tokens (as-is extraction) to JSON."""
571
+ result = {
572
+ "metadata": {
573
+ "source_url": state.base_url,
574
+ "extracted_at": datetime.now().isoformat(),
575
+ "version": "v1-stage1-extracted",
576
+ "stage": "extraction",
577
+ },
578
+ "colors": {}, # Viewport-agnostic
579
+ "typography": {
580
+ "desktop": {},
581
+ "mobile": {},
582
+ },
583
+ "spacing": {
584
+ "desktop": {},
585
+ "mobile": {},
586
+ },
587
+ "radius": {}, # Viewport-agnostic
588
+ }
589
+
590
+ # Colors (no viewport prefix - same across devices)
591
+ if state.desktop_normalized:
592
+ for name, c in state.desktop_normalized.colors.items():
593
+ result["colors"][c.suggested_name or c.value] = {
594
+ "value": c.value,
595
+ "frequency": c.frequency,
596
+ "confidence": c.confidence.value if c.confidence else "medium",
597
+ "contexts": c.contexts[:3],
598
+ }
599
+
600
+ # Typography (viewport-specific)
601
+ if state.desktop_normalized:
602
+ for name, t in state.desktop_normalized.typography.items():
603
+ key = t.suggested_name or f"{t.font_family}-{t.font_size}"
604
+ result["typography"]["desktop"][key] = {
605
+ "font_family": t.font_family,
606
+ "font_size": t.font_size,
607
+ "font_weight": t.font_weight,
608
+ "line_height": t.line_height,
609
+ "frequency": t.frequency,
610
+ }
611
+
612
+ if state.mobile_normalized:
613
+ for name, t in state.mobile_normalized.typography.items():
614
+ key = t.suggested_name or f"{t.font_family}-{t.font_size}"
615
+ result["typography"]["mobile"][key] = {
616
+ "font_family": t.font_family,
617
+ "font_size": t.font_size,
618
+ "font_weight": t.font_weight,
619
+ "line_height": t.line_height,
620
+ "frequency": t.frequency,
621
+ }
622
+
623
+ # Spacing (viewport-specific if different)
624
+ if state.desktop_normalized:
625
+ for name, s in state.desktop_normalized.spacing.items():
626
+ key = s.suggested_name or s.value
627
+ result["spacing"]["desktop"][key] = {
628
+ "value": s.value,
629
+ "value_px": s.value_px,
630
+ "fits_base_8": s.fits_base_8,
631
+ "frequency": s.frequency,
632
+ }
633
+
634
+ if state.mobile_normalized:
635
+ for name, s in state.mobile_normalized.spacing.items():
636
+ key = s.suggested_name or s.value
637
+ result["spacing"]["mobile"][key] = {
638
+ "value": s.value,
639
+ "value_px": s.value_px,
640
+ "fits_base_8": s.fits_base_8,
641
+ "frequency": s.frequency,
642
+ }
643
+
644
+ # Radius (no viewport prefix)
645
+ if state.desktop_normalized:
646
+ for name, r in state.desktop_normalized.radius.items():
647
+ result["radius"][name] = {
648
+ "value": r.value,
649
+ "frequency": r.frequency,
650
+ }
651
+
652
+ return json.dumps(result, indent=2, default=str)
653
+
654
+
655
+ def export_tokens_json():
656
+ """Export tokens to JSON."""
657
+ result = {
658
+ "metadata": {
659
+ "source_url": state.base_url,
660
+ "extracted_at": datetime.now().isoformat(),
661
+ "version": "v1-extracted",
662
+ },
663
+ "desktop": None,
664
+ "mobile": None,
665
+ }
666
+
667
+ if state.desktop_normalized:
668
+ result["desktop"] = {
669
+ "colors": [
670
+ {"value": c.value, "name": c.suggested_name, "frequency": c.frequency,
671
+ "confidence": c.confidence.value if c.confidence else "medium"}
672
+ for c in state.desktop_normalized.colors
673
+ ],
674
+ "typography": [
675
+ {"font_family": t.font_family, "font_size": t.font_size,
676
+ "font_weight": t.font_weight, "line_height": t.line_height,
677
+ "name": t.suggested_name, "frequency": t.frequency}
678
+ for t in state.desktop_normalized.typography
679
+ ],
680
+ "spacing": [
681
+ {"value": s.value, "value_px": s.value_px, "name": s.suggested_name,
682
+ "frequency": s.frequency, "fits_base_8": s.fits_base_8}
683
+ for s in state.desktop_normalized.spacing
684
+ ],
685
+ }
686
+
687
+ if state.mobile_normalized:
688
+ result["mobile"] = {
689
+ "colors": [
690
+ {"value": c.value, "name": c.suggested_name, "frequency": c.frequency,
691
+ "confidence": c.confidence.value if c.confidence else "medium"}
692
+ for c in state.mobile_normalized.colors
693
+ ],
694
+ "typography": [
695
+ {"font_family": t.font_family, "font_size": t.font_size,
696
+ "font_weight": t.font_weight, "line_height": t.line_height,
697
+ "name": t.suggested_name, "frequency": t.frequency}
698
+ for t in state.mobile_normalized.typography
699
+ ],
700
+ "spacing": [
701
+ {"value": s.value, "value_px": s.value_px, "name": s.suggested_name,
702
+ "frequency": s.frequency, "fits_base_8": s.fits_base_8}
703
+ for s in state.mobile_normalized.spacing
704
+ ],
705
+ }
706
+
707
+ return json.dumps(result, indent=2, default=str)
708
+
709
+
710
+ # =============================================================================
711
+ # UI BUILDING
712
+ # =============================================================================
713
+
714
+ def create_ui():
715
+ """Create the Gradio interface."""
716
+
717
+ with gr.Blocks(
718
+ title="Design System Extractor v2",
719
+ theme=gr.themes.Soft(),
720
+ css="""
721
+ .color-swatch { display: inline-block; width: 24px; height: 24px; border-radius: 4px; margin-right: 8px; vertical-align: middle; }
722
+ """
723
+ ) as app:
724
+
725
+ gr.Markdown("""
726
+ # 🎨 Design System Extractor v2
727
+
728
+ **Reverse-engineer design systems from live websites.**
729
+
730
+ A semi-automated, human-in-the-loop system that extracts, normalizes, and upgrades design tokens.
731
+
732
+ ---
733
+ """)
734
+
735
+ # =================================================================
736
+ # CONFIGURATION
737
+ # =================================================================
738
+
739
+ with gr.Accordion("βš™οΈ Configuration", open=not bool(HF_TOKEN_FROM_ENV)):
740
+ gr.Markdown("**HuggingFace Token** β€” Required for Stage 2 (AI upgrades)")
741
+ with gr.Row():
742
+ hf_token_input = gr.Textbox(
743
+ label="HF Token", placeholder="hf_xxxx", type="password",
744
+ scale=4, value=HF_TOKEN_FROM_ENV,
745
+ )
746
+ save_token_btn = gr.Button("πŸ’Ύ Save", scale=1)
747
+ token_status = gr.Markdown("βœ… Token loaded" if HF_TOKEN_FROM_ENV else "⏳ Enter token")
748
+
749
+ def save_token(token):
750
+ if token and len(token) > 10:
751
+ os.environ["HF_TOKEN"] = token.strip()
752
+ return "βœ… Token saved!"
753
+ return "❌ Invalid token"
754
+
755
+ save_token_btn.click(save_token, [hf_token_input], [token_status])
756
+
757
+ # =================================================================
758
+ # URL INPUT & PAGE DISCOVERY
759
+ # =================================================================
760
+
761
+ with gr.Accordion("πŸ” Step 1: Discover Pages", open=True):
762
+ gr.Markdown("Enter your website URL to discover pages for extraction.")
763
+
764
+ with gr.Row():
765
+ url_input = gr.Textbox(label="Website URL", placeholder="https://example.com", scale=4)
766
+ discover_btn = gr.Button("πŸ” Discover Pages", variant="primary", scale=1)
767
+
768
+ discover_status = gr.Markdown("")
769
+
770
+ with gr.Row():
771
+ log_output = gr.Textbox(label="πŸ“‹ Log", lines=8, interactive=False)
772
+
773
+ pages_table = gr.Dataframe(
774
+ headers=["Select", "URL", "Title", "Type", "Status"],
775
+ datatype=["bool", "str", "str", "str", "str"],
776
+ label="Discovered Pages",
777
+ interactive=True,
778
+ visible=False,
779
+ )
780
+
781
+ extract_btn = gr.Button("πŸš€ Extract Tokens (Desktop + Mobile)", variant="primary", visible=False)
782
+
783
+ # =================================================================
784
+ # STAGE 1: EXTRACTION REVIEW
785
+ # =================================================================
786
+
787
+ with gr.Accordion("πŸ“Š Stage 1: Review Extracted Tokens", open=False) as stage1_accordion:
788
+
789
+ extraction_status = gr.Markdown("")
790
+
791
+ gr.Markdown("""
792
+ **Review the extracted tokens.** Toggle between Desktop and Mobile viewports.
793
+ Accept or reject tokens, then proceed to Stage 2 for AI-powered upgrades.
794
+ """)
795
+
796
+ viewport_toggle = gr.Radio(
797
+ choices=["Desktop (1440px)", "Mobile (375px)"],
798
+ value="Desktop (1440px)",
799
+ label="Viewport",
800
+ )
801
+
802
+ with gr.Tabs():
803
+ with gr.Tab("🎨 Colors"):
804
+ colors_table = gr.Dataframe(
805
+ headers=["Accept", "Color", "Suggested Name", "Frequency", "Confidence", "Contrast", "AA", "Context"],
806
+ datatype=["bool", "str", "str", "number", "str", "str", "str", "str"],
807
+ label="Colors",
808
+ interactive=True,
809
+ )
810
+
811
+ with gr.Tab("πŸ“ Typography"):
812
+ typography_table = gr.Dataframe(
813
+ headers=["Accept", "Font", "Size", "Weight", "Line Height", "Suggested Name", "Frequency", "Confidence"],
814
+ datatype=["bool", "str", "str", "str", "str", "str", "number", "str"],
815
+ label="Typography",
816
+ interactive=True,
817
+ )
818
+
819
+ with gr.Tab("πŸ“ Spacing"):
820
+ spacing_table = gr.Dataframe(
821
+ headers=["Accept", "Value", "Pixels", "Suggested Name", "Frequency", "Base 8", "Confidence"],
822
+ datatype=["bool", "str", "str", "str", "number", "str", "str"],
823
+ label="Spacing",
824
+ interactive=True,
825
+ )
826
+
827
+ with gr.Tab("πŸ”˜ Radius"):
828
+ radius_table = gr.Dataframe(
829
+ headers=["Accept", "Value", "Frequency", "Context"],
830
+ datatype=["bool", "str", "number", "str"],
831
+ label="Border Radius",
832
+ interactive=True,
833
+ )
834
+
835
+ with gr.Row():
836
+ proceed_stage2_btn = gr.Button("➑️ Proceed to Stage 2: AI Upgrades", variant="primary")
837
+ download_stage1_btn = gr.Button("πŸ“₯ Download Stage 1 JSON", variant="secondary")
838
+
839
+ # =================================================================
840
+ # STAGE 2: AI UPGRADES
841
+ # =================================================================
842
+
843
+ with gr.Accordion("🧠 Stage 2: AI-Powered Upgrades", open=False) as stage2_accordion:
844
+
845
+ stage2_status = gr.Markdown("Click 'Analyze' to start AI-powered design system analysis.")
846
+
847
+ analyze_btn = gr.Button("πŸ€– Analyze Design System", variant="primary")
848
+
849
+ with gr.Accordion("πŸ“‹ AI Analysis Log", open=False):
850
+ stage2_log = gr.Textbox(label="Log", lines=8, interactive=False)
851
+
852
+ # =============================================================
853
+ # TYPOGRAPHY SECTION
854
+ # =============================================================
855
+ gr.Markdown("---")
856
+ gr.Markdown("## πŸ“ Typography")
857
+
858
+ with gr.Row():
859
+ with gr.Column(scale=2):
860
+ gr.Markdown("### Current vs Recommended Type Scales")
861
+ typography_comparison = gr.Dataframe(
862
+ headers=["Element", "Current", "Scale 1.2", "Scale 1.25 ⭐", "Scale 1.333"],
863
+ datatype=["str", "str", "str", "str", "str"],
864
+ label="Typography Comparison",
865
+ interactive=False,
866
+ )
867
+
868
+ with gr.Column(scale=1):
869
+ gr.Markdown("### Select Option")
870
+ type_scale_radio = gr.Radio(
871
+ choices=["Keep Current", "Scale 1.2 (Minor Third)", "Scale 1.25 (Major Third) ⭐", "Scale 1.333 (Perfect Fourth)"],
872
+ value="Keep Current",
873
+ label="Type Scale",
874
+ interactive=True,
875
+ )
876
+ gr.Markdown("*Font family will be preserved*")
877
+
878
+ # =============================================================
879
+ # COLORS SECTION
880
+ # =============================================================
881
+ gr.Markdown("---")
882
+ gr.Markdown("## 🎨 Colors")
883
+
884
+ gr.Markdown("Generate full color ramps (50-900 shades) from detected base colors:")
885
+ color_ramps_display = gr.Markdown("")
886
+
887
+ color_ramps_checkbox = gr.Checkbox(
888
+ label="βœ“ Generate color ramps (base colors preserved, adds tints/shades)",
889
+ value=True,
890
+ )
891
+
892
+ # =============================================================
893
+ # SPACING SECTION
894
+ # =============================================================
895
+ gr.Markdown("---")
896
+ gr.Markdown("## πŸ“ Spacing")
897
+
898
+ with gr.Row():
899
+ with gr.Column(scale=2):
900
+ gr.Markdown("### Current vs Grid-Aligned")
901
+ spacing_comparison = gr.Dataframe(
902
+ headers=["Current", "8px Grid", "4px Grid"],
903
+ datatype=["str", "str", "str"],
904
+ label="Spacing Comparison",
905
+ interactive=False,
906
+ )
907
+
908
+ with gr.Column(scale=1):
909
+ gr.Markdown("### Select Option")
910
+ spacing_radio = gr.Radio(
911
+ choices=["Keep Current", "8px Base Grid ⭐", "4px Base Grid"],
912
+ value="Keep Current",
913
+ label="Spacing System",
914
+ interactive=True,
915
+ )
916
+
917
+ # =============================================================
918
+ # RADIUS SECTION
919
+ # =============================================================
920
+ gr.Markdown("---")
921
+ gr.Markdown("## πŸ”˜ Border Radius")
922
+
923
+ radius_display = gr.Markdown("*Radius values detected. No changes recommended.*")
924
+
925
+ # =============================================================
926
+ # SHADOWS SECTION
927
+ # =============================================================
928
+ gr.Markdown("---")
929
+ gr.Markdown("## 🌫️ Shadows")
930
+
931
+ shadows_display = gr.Markdown("*Shadow values detected. No changes recommended.*")
932
+
933
+ # =============================================================
934
+ # APPLY SECTION
935
+ # =============================================================
936
+ gr.Markdown("---")
937
+
938
+ with gr.Row():
939
+ apply_upgrades_btn = gr.Button("✨ Apply Selected Upgrades", variant="primary", scale=2)
940
+ reset_btn = gr.Button("↩️ Reset to Original", variant="secondary", scale=1)
941
+
942
+ apply_status = gr.Markdown("")
943
+
944
+ # =================================================================
945
+ # STAGE 3: EXPORT
946
+ # =================================================================
947
+
948
+ with gr.Accordion("πŸ“¦ Stage 3: Export", open=False):
949
+ gr.Markdown("""
950
+ Export your design tokens to JSON (compatible with Figma Tokens Studio).
951
+
952
+ - **Stage 1 JSON**: Raw extracted tokens (as-is)
953
+ - **Final JSON**: Upgraded tokens with selected improvements
954
+ """)
955
+
956
+ with gr.Row():
957
+ export_stage1_btn = gr.Button("πŸ“₯ Export Stage 1 (As-Is)", variant="secondary")
958
+ export_final_btn = gr.Button("πŸ“₯ Export Final (Upgraded)", variant="primary")
959
+
960
+ export_output = gr.Code(label="Tokens JSON", language="json", lines=25)
961
+
962
+ export_stage1_btn.click(export_stage1_json, outputs=[export_output])
963
+ export_final_btn.click(export_tokens_json, outputs=[export_output])
964
+
965
+ # =================================================================
966
+ # EVENT HANDLERS
967
+ # =================================================================
968
+
969
+ # Store data for viewport toggle
970
+ desktop_data = gr.State({})
971
+ mobile_data = gr.State({})
972
+
973
+ # Discover pages
974
+ discover_btn.click(
975
+ fn=discover_pages,
976
+ inputs=[url_input],
977
+ outputs=[discover_status, log_output, pages_table],
978
+ ).then(
979
+ fn=lambda: (gr.update(visible=True), gr.update(visible=True)),
980
+ outputs=[pages_table, extract_btn],
981
+ )
982
+
983
+ # Extract tokens
984
+ extract_btn.click(
985
+ fn=extract_tokens,
986
+ inputs=[pages_table],
987
+ outputs=[extraction_status, log_output, desktop_data, mobile_data],
988
+ ).then(
989
+ fn=lambda d: (d.get("colors", []), d.get("typography", []), d.get("spacing", [])),
990
+ inputs=[desktop_data],
991
+ outputs=[colors_table, typography_table, spacing_table],
992
+ ).then(
993
+ fn=lambda: gr.update(open=True),
994
+ outputs=[stage1_accordion],
995
+ )
996
+
997
+ # Viewport toggle
998
+ viewport_toggle.change(
999
+ fn=switch_viewport,
1000
+ inputs=[viewport_toggle],
1001
+ outputs=[colors_table, typography_table, spacing_table],
1002
+ )
1003
+
1004
+ # Stage 2: Analyze
1005
+ analyze_btn.click(
1006
+ fn=run_stage2_analysis,
1007
+ outputs=[stage2_status, stage2_log, typography_comparison, spacing_comparison,
1008
+ color_ramps_display, radius_display, shadows_display],
1009
+ )
1010
+
1011
+ # Stage 2: Apply upgrades
1012
+ apply_upgrades_btn.click(
1013
+ fn=apply_selected_upgrades,
1014
+ inputs=[type_scale_radio, spacing_radio, color_ramps_checkbox],
1015
+ outputs=[apply_status, stage2_log],
1016
+ )
1017
+
1018
+ # Proceed to Stage 2 button
1019
+ proceed_stage2_btn.click(
1020
+ fn=lambda: gr.update(open=True),
1021
+ outputs=[stage2_accordion],
1022
+ )
1023
+
1024
+ # =================================================================
1025
+ # FOOTER
1026
+ # =================================================================
1027
+
1028
+ gr.Markdown("""
1029
+ ---
1030
+ **Design System Extractor v2** | Built with Playwright + Gradio + LangGraph + HuggingFace
1031
+
1032
+ *A semi-automated co-pilot for design system recovery and modernization.*
1033
+ """)
1034
+
1035
+ return app
1036
+
1037
+
1038
+ # =============================================================================
1039
+ # MAIN
1040
+ # =============================================================================
1041
+
1042
+ if __name__ == "__main__":
1043
+ app = create_ui()
1044
+ app.launch(server_name="0.0.0.0", server_port=7860)