Upload 7 files

components/AIBuilderModal.tsx

@@ -0,0 +1,184 @@
+import React, { useState, useEffect, useRef } from 'react';
+import { Sparkles, X, MessageSquare, Play, RefreshCw, AlertCircle, HardHat, Search, Wand2, CheckCircle2 } from 'lucide-react';
+import { generateGraphWithAgents, AgentStatus } from '../services/geminiService';
+import { Node } from 'reactflow';
+import { NodeData } from '../types';
+
+interface AIBuilderModalProps {
+  isOpen: boolean;
+  onClose: () => void;
+  onApply: (nodes: any[], edges: any[]) => void;
+  currentNodes: Node<NodeData>[];
+}
+
+const AIBuilderModal: React.FC<AIBuilderModalProps> = ({ isOpen, onClose, onApply, currentNodes }) => {
+  const [prompt, setPrompt] = useState('');
+  const [isBuilding, setIsBuilding] = useState(false);
+  const [error, setError] = useState<string | null>(null);
+  
+  const [agentStatus, setAgentStatus] = useState<AgentStatus>('idle');
+  const [agentMessage, setAgentMessage] = useState('');
+
+  // Reset state when opened
+  useEffect(() => {
+    if (isOpen) {
+      setAgentStatus('idle');
+      setAgentMessage('');
+      setError(null);
+    }
+  }, [isOpen]);
+
+  if (!isOpen) return null;
+
+  const handleBuild = async () => {
+    if (!prompt.trim()) return;
+    setIsBuilding(true);
+    setError(null);
+
+    try {
+      const result = await generateGraphWithAgents(prompt, currentNodes, (status, msg) => {
+          setAgentStatus(status);
+          setAgentMessage(msg);
+      });
+      
+      if (result && result.nodes && result.edges) {
+        // Ensure nodes have correct style defaults if missing
+        const processedNodes = result.nodes.map((n: any) => ({
+          ...n,
+          type: 'custom', // enforce custom type
+          data: {
+            ...n.data,
+            // ensure label exists if AI forgot it
+            label: n.data.label || n.data.type 
+          }
+        }));
+        
+        // Ensure edges have styling
+        const processedEdges = result.edges.map((e: any) => ({
+          ...e,
+          animated: true, 
+          style: { stroke: '#94a3b8' }
+        }));
+
+        // Small delay to let user see "Complete" state
+        setTimeout(() => {
+             onApply(processedNodes, processedEdges);
+             onClose();
+             setIsBuilding(false);
+        }, 1000);
+      } else {
+        throw new Error("Invalid response format from AI");
+      }
+    } catch (err) {
+      setError(err instanceof Error ? err.message : "Failed to generate architecture");
+      setAgentStatus('error');
+      setIsBuilding(false);
+    }
+  };
+
+  const getStepStatus = (step: AgentStatus, current: AgentStatus) => {
+    const order = ['idle', 'architect', 'critic', 'refiner', 'complete'];
+    const stepIdx = order.indexOf(step);
+    const currentIdx = order.indexOf(current);
+    
+    if (current === 'error') return 'text-slate-500';
+    if (currentIdx > stepIdx) return 'text-emerald-400';
+    if (currentIdx === stepIdx) return 'text-blue-400 animate-pulse';
+    return 'text-slate-600';
+  };
+
+  const renderAgentStep = (step: AgentStatus, icon: React.ReactNode, label: string) => {
+      const statusColor = getStepStatus(step, agentStatus);
+      const isCurrent = agentStatus === step;
+      const isDone = ['architect', 'critic', 'refiner', 'complete'].indexOf(agentStatus) > ['architect', 'critic', 'refiner'].indexOf(step);
+
+      return (
+        <div className={`flex items-center gap-3 p-3 rounded-lg border border-slate-800 ${isCurrent ? 'bg-slate-800/50 border-blue-500/30' : 'bg-slate-900'}`}>
+            <div className={`${statusColor} transition-colors duration-300`}>
+                {isDone ? <CheckCircle2 size={24} /> : icon}
+            </div>
+            <div className="flex-1">
+                <div className={`font-semibold text-sm ${isCurrent ? 'text-blue-200' : 'text-slate-400'}`}>{label}</div>
+                {isCurrent && (
+                     <div className="text-xs text-slate-500 mt-1">{agentMessage}</div>
+                )}
+            </div>
+            {isCurrent && <div className="w-2 h-2 rounded-full bg-blue-400 animate-ping" />}
+        </div>
+      );
+  }
+
+  return (
+    <div className="absolute inset-0 z-50 flex items-center justify-center bg-black/70 backdrop-blur-sm p-4">
+      <div className="bg-slate-950 w-full max-w-lg rounded-xl border border-slate-700 shadow-2xl flex flex-col overflow-hidden animate-in fade-in zoom-in duration-200">
+        
+        <div className="flex items-center justify-between px-6 py-4 border-b border-slate-800 bg-gradient-to-r from-slate-900 to-slate-800">
+          <h2 className="text-lg font-bold text-white flex items-center gap-2">
+            <Sparkles className="text-purple-400" size={20} />
+            AI Architect Team
+          </h2>
+          <button onClick={onClose} disabled={isBuilding} className="text-slate-400 hover:text-white transition-colors disabled:opacity-50">
+            <X size={20} />
+          </button>
+        </div>
+
+        <div className="p-6 space-y-4">
+          {!isBuilding && agentStatus !== 'complete' ? (
+            <>
+                <div className="bg-purple-500/10 border border-purple-500/20 rounded-lg p-3 text-sm text-purple-200 flex gap-3">
+                    <MessageSquare size={18} className="shrink-0 mt-0.5" />
+                    <p>
+                    Describe your model. Our AI Agents (Architect, Critic, and Refiner) will collaborate to build the best architecture for you.
+                    </p>
+                </div>
+
+                <textarea
+                    className="w-full h-32 bg-slate-900 border border-slate-700 rounded-lg p-3 text-slate-200 focus:ring-2 focus:ring-purple-500 outline-none resize-none placeholder-slate-600"
+                    placeholder="e.g. Build a robust CNN for medical image segmentation with attention gates..."
+                    value={prompt}
+                    onChange={(e) => setPrompt(e.target.value)}
+                />
+
+                {error && (
+                    <div className="flex items-center gap-2 text-red-400 text-sm bg-red-500/10 p-2 rounded border border-red-500/20">
+                    <AlertCircle size={14} />
+                    {error}
+                    </div>
+                )}
+
+                <div className="flex justify-end pt-2">
+                    <button
+                    onClick={handleBuild}
+                    disabled={!prompt.trim()}
+                    className={`
+                        flex items-center gap-2 px-6 py-2 rounded-lg font-medium text-white transition-all
+                        ${!prompt.trim() 
+                        ? 'bg-slate-800 cursor-not-allowed opacity-50' 
+                        : 'bg-gradient-to-r from-purple-600 to-blue-600 hover:from-purple-500 hover:to-blue-500 shadow-lg shadow-purple-500/20'}
+                    `}
+                    >
+                    <Play size={18} fill="currentColor" />
+                    Start Agents
+                    </button>
+                </div>
+            </>
+          ) : (
+            <div className="space-y-3 py-2">
+                {renderAgentStep('architect', <HardHat size={24} />, "Architect Agent: Drafting Layout")}
+                {renderAgentStep('critic', <Search size={24} />, "Critic Agent: Reviewing Design")}
+                {renderAgentStep('refiner', <Wand2 size={24} />, "Refiner Agent: Finalizing Architecture")}
+                
+                {agentStatus === 'complete' && (
+                     <div className="text-center text-emerald-400 font-bold mt-4 animate-pulse">
+                        Architecture Generation Complete!
+                     </div>
+                )}
+            </div>
+          )}
+        </div>
+      </div>
+    </div>
+  );
+};
+
+export default AIBuilderModal;

components/ApiKeyModal.tsx

@@ -0,0 +1,84 @@
+import React, { useState } from 'react';
+import { Key, X, Check, Lock } from 'lucide-react';
+import { setUserApiKey } from '../services/geminiService';
+
+interface ApiKeyModalProps {
+  isOpen: boolean;
+  onClose: () => void;
+  onSuccess: () => void;
+}
+
+const ApiKeyModal: React.FC<ApiKeyModalProps> = ({ isOpen, onClose, onSuccess }) => {
+  const [key, setKey] = useState('');
+  const [saved, setSaved] = useState(false);
+
+  if (!isOpen) return null;
+
+  const handleSave = () => {
+    if (key.trim()) {
+      setUserApiKey(key.trim());
+      setSaved(true);
+      setTimeout(() => {
+        onSuccess();
+        onClose();
+        setSaved(false);
+      }, 1000);
+    }
+  };
+
+  return (
+    <div className="absolute inset-0 z-50 flex items-center justify-center bg-black/70 backdrop-blur-sm p-4">
+      <div className="bg-slate-900 w-full max-w-md rounded-xl border border-slate-700 shadow-2xl overflow-hidden animate-in fade-in zoom-in duration-200">
+        <div className="flex items-center justify-between px-6 py-4 border-b border-slate-800 bg-slate-800/50">
+          <h2 className="text-lg font-bold text-white flex items-center gap-2">
+            <Key className="text-emerald-400" size={20} />
+            API Access
+          </h2>
+          <button onClick={onClose} className="text-slate-400 hover:text-white transition-colors">
+            <X size={20} />
+          </button>
+        </div>
+        
+        <div className="p-6 space-y-4">
+          <p className="text-sm text-slate-400">
+            Enter your Google Gemini API Key to enable AI features like Agents, Suggestions, and Code Generation.
+          </p>
+          
+          <div className="relative">
+            <Lock className="absolute left-3 top-2.5 text-slate-500" size={16} />
+            <input 
+              type="password" 
+              className="w-full bg-slate-950 border border-slate-700 rounded-lg pl-10 pr-4 py-2 text-slate-200 focus:ring-2 focus:ring-emerald-500 outline-none placeholder-slate-600 font-mono text-sm"
+              placeholder="Paste your API key here..."
+              value={key}
+              onChange={(e) => setKey(e.target.value)}
+            />
+          </div>
+
+          <div className="bg-emerald-500/10 border border-emerald-500/20 rounded-lg p-3">
+             <p className="text-xs text-emerald-200">
+               Your key is stored locally in your browser and used directly with Google's API.
+             </p>
+          </div>
+
+          <div className="flex justify-end pt-2">
+            <button 
+              onClick={handleSave}
+              disabled={!key.trim()}
+              className={`
+                flex items-center gap-2 px-6 py-2 rounded-lg font-medium text-white transition-all
+                ${saved ? 'bg-emerald-600' : 'bg-emerald-600 hover:bg-emerald-500'}
+                ${!key.trim() && !saved ? 'opacity-50 cursor-not-allowed bg-slate-700' : ''}
+              `}
+            >
+              {saved ? <Check size={18} /> : <Key size={18} />}
+              {saved ? 'Saved!' : 'Save Key'}
+            </button>
+          </div>
+        </div>
+      </div>
+    </div>
+  );
+};
+
+export default ApiKeyModal;

components/CodeViewer.tsx

@@ -0,0 +1,188 @@
+import React, { useState, useEffect } from 'react';
+import { X, Copy, Check, Terminal, HardHat, Search, Wand2, CheckCircle2, ArrowLeft, Play } from 'lucide-react';
+import { generateCodeWithAgents, AgentStatus } from '../services/geminiService';
+
+interface CodeViewerProps {
+  code: string; // This represents the PROMPT text initially
+  isOpen: boolean;
+  onClose: () => void;
+  isLoading: boolean;
+}
+
+const CodeViewer: React.FC<CodeViewerProps> = ({ code, isOpen, onClose, isLoading }) => {
+  const [copied, setCopied] = useState(false);
+  const [viewMode, setViewMode] = useState<'prompt' | 'building' | 'code'>('prompt');
+  
+  const [finalCode, setFinalCode] = useState('');
+  const [agentStatus, setAgentStatus] = useState<AgentStatus>('idle');
+  const [agentMessage, setAgentMessage] = useState('');
+
+  // Reset state when opening/closing or changing input
+  useEffect(() => {
+    if (isOpen) {
+      setViewMode('prompt');
+      setFinalCode('');
+      setAgentStatus('idle');
+    }
+  }, [isOpen, code]);
+
+  const handleCopy = (text: string) => {
+    navigator.clipboard.writeText(text);
+    setCopied(true);
+    setTimeout(() => setCopied(false), 2000);
+  };
+
+  const handleBuildWithAgents = async () => {
+    setViewMode('building');
+    setAgentStatus('architect');
+    setAgentMessage('Initializing agents...');
+
+    try {
+        const result = await generateCodeWithAgents(code, (status, msg) => {
+            setAgentStatus(status);
+            setAgentMessage(msg);
+        });
+        setFinalCode(result);
+        // Small delay to show completion state
+        setTimeout(() => {
+            setViewMode('code');
+        }, 1000);
+    } catch (error) {
+        setAgentStatus('error');
+        setAgentMessage('Generation failed. Please try again.');
+        setTimeout(() => setViewMode('prompt'), 2000);
+    }
+  };
+
+  const getStepStatus = (step: AgentStatus, current: AgentStatus) => {
+    const order = ['idle', 'architect', 'critic', 'refiner', 'complete'];
+    const stepIdx = order.indexOf(step);
+    const currentIdx = order.indexOf(current);
+    
+    if (current === 'error') return 'text-slate-500';
+    if (currentIdx > stepIdx) return 'text-emerald-400';
+    if (currentIdx === stepIdx) return 'text-blue-400 animate-pulse';
+    return 'text-slate-600';
+  };
+
+  const renderAgentStep = (step: AgentStatus, icon: React.ReactNode, label: string) => {
+      const statusColor = getStepStatus(step, agentStatus);
+      const isCurrent = agentStatus === step;
+      const isDone = ['architect', 'critic', 'refiner', 'complete'].indexOf(agentStatus) > ['architect', 'critic', 'refiner'].indexOf(step);
+
+      return (
+        <div className={`flex items-center gap-3 p-4 rounded-lg border border-slate-800 ${isCurrent ? 'bg-slate-800/50 border-blue-500/30' : 'bg-slate-900'}`}>
+            <div className={`${statusColor} transition-colors duration-300`}>
+                {isDone ? <CheckCircle2 size={24} /> : icon}
+            </div>
+            <div className="flex-1">
+                <div className={`font-semibold text-sm ${isCurrent ? 'text-blue-200' : 'text-slate-400'}`}>{label}</div>
+                {isCurrent && (
+                     <div className="text-xs text-slate-500 mt-1">{agentMessage}</div>
+                )}
+            </div>
+            {isCurrent && <div className="w-2 h-2 rounded-full bg-blue-400 animate-ping" />}
+        </div>
+      );
+  };
+
+  if (!isOpen) return null;
+
+  return (
+    <div className="absolute inset-0 z-50 flex items-center justify-center bg-black/70 backdrop-blur-sm p-4 md:p-8">
+      <div className="bg-slate-900 w-full max-w-5xl h-[85vh] rounded-xl border border-slate-700 shadow-2xl flex flex-col overflow-hidden animate-in fade-in zoom-in duration-200">
+        
+        {/* Header */}
+        <div className="flex items-center justify-between px-6 py-4 border-b border-slate-800 bg-slate-800/50">
+          <h2 className="text-lg font-bold text-white flex items-center gap-2">
+            {viewMode === 'prompt' && <span className="text-blue-400">Architecture Prompt</span>}
+            {viewMode === 'building' && <span className="text-purple-400 flex items-center gap-2"><HardHat size={18}/> Agents Building...</span>}
+            {viewMode === 'code' && <span className="text-emerald-400 flex items-center gap-2"><Terminal size={18}/> Final Code</span>}
+            
+            {isLoading && <span className="text-xs text-blue-400 animate-pulse ml-2">Building...</span>}
+          </h2>
+          
+          <div className="flex items-center gap-3">
+             {viewMode === 'code' && (
+               <button 
+                  onClick={() => setViewMode('prompt')}
+                  className="flex items-center gap-2 px-3 py-1.5 bg-slate-800 hover:bg-slate-700 rounded text-xs font-medium text-slate-300 transition-colors border border-slate-700"
+                >
+                  <ArrowLeft size={14} /> Back to Prompt
+                </button>
+             )}
+
+            {viewMode !== 'building' && (
+              <button 
+                onClick={() => handleCopy(viewMode === 'code' ? finalCode : code)}
+                className="flex items-center gap-2 px-3 py-1.5 bg-slate-800 hover:bg-slate-700 rounded text-xs font-medium text-slate-300 transition-colors border border-slate-700"
+              >
+                {copied ? <Check size={14} className="text-emerald-400"/> : <Copy size={14} />}
+                {copied ? 'Copied' : 'Copy'}
+              </button>
+            )}
+            
+            {viewMode === 'prompt' && !isLoading && (
+              <button 
+                onClick={handleBuildWithAgents}
+                className="flex items-center gap-2 px-3 py-1.5 bg-purple-600 hover:bg-purple-500 rounded text-xs font-medium text-white transition-colors shadow-lg shadow-purple-500/20"
+              >
+                <Play size={14} fill="currentColor" />
+                Build with Agents (Beta)
+              </button>
+            )}
+            
+            <button onClick={onClose} className="text-slate-400 hover:text-white transition-colors ml-2">
+              <X size={20} />
+            </button>
+          </div>
+        </div>
+
+        {/* Content */}
+        <div className="flex-1 overflow-hidden bg-[#0d1117] relative">
+          
+          {/* View: Loading Prompt */}
+          {isLoading && viewMode === 'prompt' && (
+            <div className="h-full flex flex-col items-center justify-center space-y-4">
+              <div className="w-12 h-12 border-4 border-blue-500/30 border-t-blue-500 rounded-full animate-spin"></div>
+              <p className="text-slate-400 text-sm animate-pulse">Refining architecture prompt...</p>
+            </div>
+          )}
+
+          {/* View: Prompt Text */}
+          {!isLoading && viewMode === 'prompt' && (
+             <div className="h-full overflow-auto p-6 font-mono text-sm text-slate-300 leading-relaxed whitespace-pre-wrap">
+               {code}
+             </div>
+          )}
+
+          {/* View: Agents Building */}
+          {viewMode === 'building' && (
+             <div className="h-full flex flex-col items-center justify-center p-8">
+                 <div className="w-full max-w-md space-y-3">
+                    {renderAgentStep('architect', <HardHat size={24} />, "Coder Agent: Writing Implementation")}
+                    {renderAgentStep('critic', <Search size={24} />, "Reviewer Agent: Analyzing Logic")}
+                    {renderAgentStep('refiner', <Wand2 size={24} />, "Polisher Agent: Finalizing Code")}
+                 </div>
+                 {agentStatus === 'complete' && (
+                     <div className="mt-8 text-emerald-400 font-bold animate-bounce">
+                         Done! Loading code...
+                     </div>
+                 )}
+             </div>
+          )}
+
+          {/* View: Final Code */}
+          {viewMode === 'code' && (
+            <div className="h-full overflow-auto p-6 font-mono text-sm leading-relaxed whitespace-pre-wrap text-blue-100 selection:bg-blue-500/30">
+              {finalCode}
+            </div>
+          )}
+
+        </div>
+      </div>
+    </div>
+  );
+};
+
+export default CodeViewer;

components/CustomNode.tsx

@@ -0,0 +1,183 @@
+import React, { memo } from 'react';
+import { Handle, Position, NodeProps } from 'reactflow';
+import { 
+  Layers, Box, Activity, Zap, ArrowRight, Grid, Minimize, 
+  Database, GitBranch, AlignJustify, Type, Combine, Maximize,
+  ArrowUpCircle, Sliders, RefreshCcw, Brain, Crosshair, Network, Clock, Eye, Workflow,
+  Terminal, MinusCircle
+} from 'lucide-react';
+import { NodeData, LayerType } from '../types';
+
+const getIcon = (type: LayerType) => {
+  switch (type) {
+    case LayerType.INPUT: return <ArrowRight className="w-4 h-4" />;
+    case LayerType.CONV1D: return <Activity className="w-4 h-4" />;
+    case LayerType.CONV2D: return <Layers className="w-4 h-4" />;
+    case LayerType.CONV_TRANSPOSE2D: return <ArrowUpCircle className="w-4 h-4" />;
+    case LayerType.LINEAR: return <Grid className="w-4 h-4" />;
+    case LayerType.RELU: 
+    case LayerType.LEAKYRELU: 
+    case LayerType.GELU: 
+    case LayerType.SILU: 
+    case LayerType.SWIGLU:
+    case LayerType.SIGMOID: 
+    case LayerType.TANH: return <Zap className="w-4 h-4" />;
+    case LayerType.MAXPOOL: 
+    case LayerType.AVGPOOL:
+    case LayerType.ADAPTIVEAVGPOOL: return <Minimize className="w-4 h-4" />;
+    case LayerType.UPSAMPLE: return <Maximize className="w-4 h-4" />;
+    
+    // Transformer / GenAI
+    case LayerType.ATTENTION: 
+    case LayerType.CROSS_ATTENTION: return <RefreshCcw className="w-4 h-4" />;
+    case LayerType.TRANSFORMER_BLOCK:
+    case LayerType.TRANSFORMER_ENCODER:
+    case LayerType.TRANSFORMER_DECODER: return <Brain className="w-4 h-4" />;
+    case LayerType.MOE_BLOCK: return <Network className="w-4 h-4" />;
+    case LayerType.ACTION_HEAD: return <Crosshair className="w-4 h-4" />;
+    case LayerType.PATCH_EMBED:
+    case LayerType.SAM_PROMPT_ENCODER: return <Eye className="w-4 h-4" />;
+    case LayerType.SAM_MASK_DECODER: return <Workflow className="w-4 h-4" />;
+    case LayerType.TIME_EMBEDDING:
+    case LayerType.ROPE: return <Clock className="w-4 h-4" />;
+    
+    case LayerType.LSTM:
+    case LayerType.GRU: return <GitBranch className="w-4 h-4" />;
+    case LayerType.EMBEDDING: return <Database className="w-4 h-4" />;
+    case LayerType.LAYERNORM:
+    case LayerType.BATCHNORM: 
+    case LayerType.INSTANCENORM: 
+    case LayerType.RMSNORM: return <Sliders className="w-4 h-4" />;
+    case LayerType.FLATTEN: return <Type className="w-4 h-4" />;
+    case LayerType.CONCAT:
+    case LayerType.ADD: return <Combine className="w-4 h-4" />;
+    
+    case LayerType.CUSTOM: return <Terminal className="w-4 h-4" />;
+    case LayerType.IDENTITY: return <MinusCircle className="w-4 h-4" />;
+    
+    default: return <Box className="w-4 h-4" />;
+  }
+};
+
+const getColor = (type: LayerType) => {
+  switch (type) {
+    case LayerType.INPUT: return 'border-emerald-500 shadow-emerald-500/20';
+    case LayerType.OUTPUT: return 'border-red-500 shadow-red-500/20';
+    
+    case LayerType.CONV2D: 
+    case LayerType.CONV1D:
+    case LayerType.CONV_TRANSPOSE2D:
+    case LayerType.MAXPOOL: 
+    case LayerType.AVGPOOL:
+    case LayerType.ADAPTIVEAVGPOOL:
+    case LayerType.UPSAMPLE: return 'border-blue-500 shadow-blue-500/20';
+    
+    case LayerType.LINEAR: 
+    case LayerType.ACTION_HEAD: return 'border-violet-500 shadow-violet-500/20';
+    
+    // GenAI / Transformer - Gold/Amber theme
+    case LayerType.ATTENTION: 
+    case LayerType.CROSS_ATTENTION:
+    case LayerType.TRANSFORMER_BLOCK:
+    case LayerType.TRANSFORMER_ENCODER:
+    case LayerType.TRANSFORMER_DECODER:
+    case LayerType.MOE_BLOCK:
+    case LayerType.PATCH_EMBED:
+    case LayerType.SAM_PROMPT_ENCODER:
+    case LayerType.SAM_MASK_DECODER: return 'border-amber-500 shadow-amber-500/20';
+    
+    case LayerType.LSTM:
+    case LayerType.GRU:
+    case LayerType.EMBEDDING: 
+    case LayerType.TIME_EMBEDDING:
+    case LayerType.ROPE: return 'border-orange-500 shadow-orange-500/20';
+    
+    case LayerType.BATCHNORM:
+    case LayerType.LAYERNORM: 
+    case LayerType.INSTANCENORM: 
+    case LayerType.RMSNORM: return 'border-cyan-500 shadow-cyan-500/20';
+    
+    case LayerType.CONCAT:
+    case LayerType.ADD: return 'border-pink-500 shadow-pink-500/20';
+    
+    case LayerType.CUSTOM: return 'border-lime-500 shadow-lime-500/20';
+    case LayerType.IDENTITY: return 'border-slate-600 shadow-slate-600/20';
+
+    default: return 'border-slate-500 shadow-slate-500/20';
+  }
+};
+
+const CustomNode = ({ data, selected }: NodeProps<NodeData>) => {
+  const isInput = data.type === LayerType.INPUT;
+  const isOutput = data.type === LayerType.OUTPUT;
+  
+  // Render input modality badge
+  const inputBadge = isInput && data.params.modality ? (
+    <span className="ml-auto text-[9px] px-1.5 py-0.5 rounded-full bg-slate-700 text-emerald-400 font-mono tracking-tighter uppercase">
+      {data.params.modality}
+    </span>
+  ) : null;
+  
+  // Custom layer name badge
+  const customBadge = data.type === LayerType.CUSTOM && data.params.class_name ? (
+     <span className="ml-auto text-[8px] px-1.5 py-0.5 rounded-full bg-slate-700 text-lime-400 font-mono truncate max-w-[80px]">
+      {data.params.class_name.split('.').pop()}
+    </span>
+  ) : null;
+
+  return (
+    <div className={`
+      relative min-w-[150px] bg-slate-900 rounded-lg border-2 
+      transition-all duration-200 shadow-lg
+      ${getColor(data.type)}
+      ${selected ? 'ring-2 ring-white ring-offset-2 ring-offset-slate-900 scale-105' : ''}
+    `}>
+      {/* Input Handle */}
+      {!isInput && (
+        <Handle 
+          type="target" 
+          position={Position.Top} 
+          className="!bg-slate-400 !w-3 !h-3 !-top-2 !border-slate-900" 
+        />
+      )}
+
+      {/* Header */}
+      <div className="flex items-center gap-2 px-3 py-2 border-b border-slate-800 bg-slate-800/50 rounded-t-md">
+        <span className="text-slate-300">
+          {getIcon(data.type)}
+        </span>
+        <span className="text-xs font-bold uppercase tracking-wider text-slate-200 truncate max-w-[100px]">
+          {data.label}
+        </span>
+        {inputBadge}
+        {customBadge}
+      </div>
+
+      {/* Body - Parameters Summary */}
+      <div className="p-3">
+        <div className="text-[10px] text-slate-400 font-mono space-y-1">
+          {Object.entries(data.params).slice(0, 3).map(([key, value]) => (
+            <div key={key} className="flex justify-between gap-4">
+              <span className="opacity-70 truncate max-w-[80px]">{key}:</span>
+              <span className="text-slate-200 truncate max-w-[80px]">{String(value)}</span>
+            </div>
+          ))}
+          {Object.keys(data.params).length === 0 && (
+             <span className="opacity-50 italic">No parameters</span>
+          )}
+        </div>
+      </div>
+
+      {/* Output Handle */}
+      {!isOutput && (
+        <Handle 
+          type="source" 
+          position={Position.Bottom} 
+          className="!bg-slate-400 !w-3 !h-3 !-bottom-2 !border-slate-900" 
+        />
+      )}
+    </div>
+  );
+};
+
+export default memo(CustomNode);

components/PropertiesPanel.tsx

@@ -0,0 +1,112 @@
+import React from 'react';
+import { Node } from 'reactflow';
+import { NodeData, LayerDefinition } from '../types';
+import { LAYER_DEFINITIONS } from '../constants';
+import { X, Trash2 } from 'lucide-react';
+
+interface PropertiesPanelProps {
+  selectedNode: Node<NodeData> | null;
+  onChange: (id: string, newParams: Record<string, any>) => void;
+  onDelete: (id: string) => void;
+  onClose: () => void;
+}
+
+const PropertiesPanel: React.FC<PropertiesPanelProps> = ({ selectedNode, onChange, onDelete, onClose }) => {
+  if (!selectedNode) {
+    return (
+      <div className="w-80 bg-slate-900 border-l border-slate-800 p-6 flex flex-col items-center justify-center text-slate-500 h-full">
+        <div className="w-16 h-16 rounded-full bg-slate-800 mb-4 flex items-center justify-center">
+          <span className="text-2xl opacity-50">⚡</span>
+        </div>
+        <p className="text-sm text-center">Select a node on the canvas to configure parameters.</p>
+      </div>
+    );
+  }
+
+  const definition: LayerDefinition = LAYER_DEFINITIONS[selectedNode.data.type];
+  
+  const handleParamChange = (name: string, value: any, type: string) => {
+    let parsedValue = value;
+    if (type === 'number') parsedValue = Number(value);
+    if (type === 'boolean') parsedValue = value === 'true';
+
+    onChange(selectedNode.id, {
+      ...selectedNode.data.params,
+      [name]: parsedValue
+    });
+  };
+
+  return (
+    <div className="w-80 bg-slate-900 border-l border-slate-800 flex flex-col h-full">
+      <div className="p-4 border-b border-slate-800 flex justify-between items-center">
+        <div>
+          <h2 className="text-lg font-bold text-slate-100">{definition.label}</h2>
+          <p className="text-xs text-slate-500 font-mono">{selectedNode.id}</p>
+        </div>
+        <button onClick={onClose} className="text-slate-500 hover:text-slate-300">
+          <X size={18} />
+        </button>
+      </div>
+
+      <div className="flex-1 overflow-y-auto p-4 space-y-5">
+        <div className="text-sm text-slate-400 italic bg-slate-800/50 p-3 rounded border border-slate-800">
+          {definition.description}
+        </div>
+
+        <div className="space-y-4">
+          {definition.parameters.map((param) => (
+            <div key={param.name} className="space-y-1.5">
+              <label className="block text-xs font-semibold text-slate-300 uppercase tracking-wide">
+                {param.label}
+              </label>
+              
+              {param.type === 'select' ? (
+                <select
+                  className="w-full bg-slate-950 border border-slate-700 rounded px-3 py-2 text-sm text-slate-200 focus:ring-1 focus:ring-blue-500 outline-none"
+                  value={selectedNode.data.params[param.name] || param.default}
+                  onChange={(e) => handleParamChange(param.name, e.target.value, param.type)}
+                >
+                  {param.options?.map(opt => (
+                    <option key={opt} value={opt}>{opt}</option>
+                  ))}
+                </select>
+              ) : param.type === 'boolean' ? (
+                 <select
+                  className="w-full bg-slate-950 border border-slate-700 rounded px-3 py-2 text-sm text-slate-200 focus:ring-1 focus:ring-blue-500 outline-none"
+                  value={String(selectedNode.data.params[param.name] ?? param.default)}
+                  onChange={(e) => handleParamChange(param.name, e.target.value === 'true', param.type)}
+                >
+                  <option value="true">True</option>
+                  <option value="false">False</option>
+                </select>
+              ) : (
+                <input
+                  type={param.type === 'number' ? 'number' : 'text'}
+                  className="w-full bg-slate-950 border border-slate-700 rounded px-3 py-2 text-sm text-slate-200 focus:ring-1 focus:ring-blue-500 outline-none placeholder-slate-700"
+                  value={selectedNode.data.params[param.name] ?? param.default}
+                  onChange={(e) => handleParamChange(param.name, e.target.value, param.type)}
+                />
+              )}
+            </div>
+          ))}
+          
+          {definition.parameters.length === 0 && (
+             <p className="text-sm text-slate-500 text-center py-4">This layer has no configurable parameters.</p>
+          )}
+        </div>
+      </div>
+
+      <div className="p-4 border-t border-slate-800">
+        <button
+          onClick={() => onDelete(selectedNode.id)}
+          className="w-full flex items-center justify-center gap-2 bg-red-500/10 hover:bg-red-500/20 text-red-500 py-2 rounded transition-colors text-sm font-medium border border-red-500/20"
+        >
+          <Trash2 size={16} />
+          Delete Node
+        </button>
+      </div>
+    </div>
+  );
+};
+
+export default PropertiesPanel;

components/Sidebar.tsx

@@ -0,0 +1,89 @@
+import React from 'react';
+import { LAYER_DEFINITIONS } from '../constants';
+import { LayerType } from '../types';
+import { Box, Sparkles, LayoutTemplate, Circle } from 'lucide-react';
+
+interface SidebarProps {
+  onOpenAIBuilder: () => void;
+  onSelectTemplate: (templateId: string) => void;
+  isConnected: boolean;
+}
+
+const Sidebar: React.FC<SidebarProps> = ({ onOpenAIBuilder, onSelectTemplate, isConnected }) => {
+  const onDragStart = (event: React.DragEvent, layerType: LayerType) => {
+    event.dataTransfer.setData('application/reactflow', layerType);
+    event.dataTransfer.effectAllowed = 'move';
+  };
+
+  const categories = Array.from(new Set(Object.values(LAYER_DEFINITIONS).map(l => l.category)));
+  const categoryOrder = ['Core', 'Convolution', 'Recurrent', 'Transformer', 'Normalization', 'Utility'];
+  categories.sort((a, b) => categoryOrder.indexOf(a) - categoryOrder.indexOf(b));
+
+  return (
+    <aside className="w-64 bg-slate-900 border-r border-slate-800 flex flex-col h-full z-10">
+      <div className="p-4 border-b border-slate-800 bg-slate-900">
+        <h1 className="text-xl font-bold bg-gradient-to-r from-blue-400 to-violet-400 bg-clip-text text-transparent">
+          Architecture Agents
+        </h1>
+        <p className="text-xs text-slate-500 mt-1">AI-Powered Model Builder</p>
+        
+        <div className="mt-4 grid grid-cols-2 gap-2">
+          <button 
+            onClick={onOpenAIBuilder}
+            className="flex flex-col items-center justify-center p-2 bg-purple-500/10 hover:bg-purple-500/20 border border-purple-500/30 rounded-lg text-purple-300 transition-colors group"
+          >
+            <Sparkles size={18} className="mb-1 group-hover:scale-110 transition-transform" />
+            <span className="text-[10px] font-bold">AI Builder</span>
+          </button>
+          <button 
+            onClick={() => onSelectTemplate('menu')}
+            className="flex flex-col items-center justify-center p-2 bg-blue-500/10 hover:bg-blue-500/20 border border-blue-500/30 rounded-lg text-blue-300 transition-colors group"
+          >
+            <LayoutTemplate size={18} className="mb-1 group-hover:scale-110 transition-transform" />
+            <span className="text-[10px] font-bold">Templates</span>
+          </button>
+        </div>
+      </div>
+      
+      <div className="flex-1 overflow-y-auto p-4 space-y-6 scrollbar-thin scrollbar-thumb-slate-700">
+        {categories.map(category => (
+          <div key={category}>
+            <h3 className="text-xs font-semibold text-slate-500 uppercase tracking-widest mb-3 flex items-center gap-2">
+              {category}
+              <div className="h-px flex-1 bg-slate-800"></div>
+            </h3>
+            <div className="grid grid-cols-1 gap-2">
+              {Object.values(LAYER_DEFINITIONS)
+                .filter(l => l.category === category)
+                .map(layer => (
+                  <div
+                    key={layer.type}
+                    className="bg-slate-800 hover:bg-slate-750 p-3 rounded border border-slate-700 cursor-grab active:cursor-grabbing transition-colors group relative overflow-hidden"
+                    onDragStart={(event) => onDragStart(event, layer.type)}
+                    draggable
+                  >
+                    <div className="flex items-center gap-3 relative z-10">
+                      <div className={`p-1.5 rounded transition-colors group-hover:bg-slate-900 bg-slate-900/50`}>
+                        <Box size={14} className="text-slate-400 group-hover:text-blue-400" />
+                      </div>
+                      <div>
+                        <div className="text-sm font-medium text-slate-200 group-hover:text-white">{layer.label}</div>
+                        <div className="text-xs text-slate-500 leading-tight group-hover:text-slate-400">{layer.description}</div>
+                      </div>
+                    </div>
+                  </div>
+                ))}
+            </div>
+          </div>
+        ))}
+      </div>
+      
+      <div className="p-4 border-t border-slate-800 text-[10px] text-slate-500 text-center flex items-center justify-center gap-2">
+         <span>v1.2.0 • Powered by Gemini 2.5</span>
+         <Circle size={8} className={isConnected ? "fill-emerald-500 text-emerald-500" : "fill-red-500 text-red-500"} />
+      </div>
+    </aside>
+  );
+};
+
+export default Sidebar;

components/SuggestionsModal.tsx

@@ -0,0 +1,78 @@
+import React from 'react';
+import { Lightbulb, X, Sparkles, Wand2 } from 'lucide-react';
+
+interface SuggestionsModalProps {
+  isOpen: boolean;
+  onClose: () => void;
+  suggestions: string;
+  isLoading: boolean;
+  onImplement: () => void;
+  isImplementing: boolean;
+}
+
+const SuggestionsModal: React.FC<SuggestionsModalProps> = ({ 
+  isOpen, 
+  onClose, 
+  suggestions, 
+  isLoading, 
+  onImplement, 
+  isImplementing 
+}) => {
+  if (!isOpen) return null;
+
+  return (
+    <div className="absolute inset-0 z-50 flex items-center justify-center bg-black/70 backdrop-blur-sm p-4">
+      <div className="bg-slate-900 w-full max-w-lg rounded-xl border border-slate-700 shadow-2xl flex flex-col overflow-hidden animate-in fade-in zoom-in duration-200">
+        <div className="flex items-center justify-between px-6 py-4 border-b border-slate-800 bg-gradient-to-r from-slate-900 to-slate-800">
+          <h2 className="text-lg font-bold text-white flex items-center gap-2">
+            <Lightbulb className="text-amber-400" size={20} />
+            AI Suggestions
+          </h2>
+          <button onClick={onClose} className="text-slate-400 hover:text-white transition-colors">
+            <X size={20} />
+          </button>
+        </div>
+        <div className="p-6 overflow-y-auto max-h-[60vh]">
+          {isLoading ? (
+            <div className="flex flex-col items-center justify-center space-y-4 py-8">
+               <Sparkles className="text-amber-400 animate-spin" size={32} />
+               <p className="text-slate-400 text-sm animate-pulse">Analyzing architecture...</p>
+            </div>
+          ) : isImplementing ? (
+            <div className="flex flex-col items-center justify-center space-y-4 py-8">
+               <Wand2 className="text-purple-400 animate-pulse" size={32} />
+               <p className="text-slate-400 text-sm animate-pulse">Agents are applying fixes...</p>
+            </div>
+          ) : (
+            <div className="text-slate-200 text-sm leading-relaxed whitespace-pre-wrap font-medium">
+              {suggestions || "No suggestions available."}
+            </div>
+          )}
+        </div>
+        <div className="p-4 border-t border-slate-800 bg-slate-900/50 flex justify-between items-center">
+           <div className="text-xs text-slate-500">
+              Review carefully before applying.
+           </div>
+           <div className="flex gap-3">
+             <button 
+               onClick={onClose} 
+               className="px-4 py-2 bg-slate-800 hover:bg-slate-700 text-slate-200 rounded text-sm transition-colors"
+             >
+               Close
+             </button>
+             {!isLoading && !isImplementing && suggestions && !suggestions.includes("Error") && (
+                <button 
+                  onClick={onImplement}
+                  className="flex items-center gap-2 px-4 py-2 bg-purple-600 hover:bg-purple-500 text-white rounded text-sm font-medium transition-colors shadow-lg shadow-purple-900/20"
+                >
+                  <Wand2 size={16} />
+                  Implement Fixes (Auto)
+                </button>
+             )}
+           </div>
+        </div>
+      </div>
+    </div>
+  );
+};
+export default SuggestionsModal;