Update components/Builder.tsx

components/Builder.tsx

@@ -12,7 +12,7 @@ import ReactFlow, {
   Panel,
   ConnectionMode
 } from 'reactflow';
-import { FileText, CheckCircle, AlertTriangle, X, Copy, Check, Lightbulb, Key, Wrench, Menu, Activity } from 'lucide-react';
+import { FileText, CheckCircle, AlertTriangle, X, Copy, Check, Lightbulb, Key, Wrench, Menu, Activity, ChevronDown, Cpu } from 'lucide-react';
 import Sidebar from './Sidebar';
 import PropertiesPanel from './PropertiesPanel';
 import CustomNode from './CustomNode';
@@ -22,9 +22,18 @@ import SuggestionsModal from './SuggestionsModal';
 import ApiKeyModal from './ApiKeyModal';
 import FixerModal from './FixerModal';
 import { INITIAL_NODES, INITIAL_EDGES, LAYER_DEFINITIONS, TEMPLATES } from '../constants';
-import { generateRefinedPrompt, validateArchitecture, getArchitectureSuggestions, getUserApiKey, implementArchitectureSuggestions } from '../services/geminiService';
+import { generateRefinedPrompt, validateArchitecture, getArchitectureSuggestions, getUserApiKey, implementArchitectureSuggestions, DEFAULT_MODEL } from '../services/geminiService';
 import { NodeData, LayerType, LogEntry } from '../types';
 
+const AVAILABLE_MODELS = [
+  { id: 'gemini-2.5-flash', name: 'Gemini 2.5 Flash', description: 'Fast & Efficient' },
+  { id: 'gemini-2.5-flash-lite', name: 'Gemini 2.5 Flash-Lite', description: 'Ultra-Fast & Efficient' },
+  { id: 'gemini-2.5-pro', name: 'Gemini 2.5 Pro', description: 'Balanced Performance' },
+  { id: 'gemini-3-flash-preview', name: 'Gemini 3 Flash', description: 'Next-Gen Speed' },
+  { id: 'gemini-3.1-flash-lite-preview', name: 'Gemini 3.1 Flash Lite', description: 'Lightweight & Fast' },
+  { id: 'gemini-3.1-pro-preview', name: 'Gemini 3.1 Pro', description: 'Complex Reasoning' },
+];
+
 let id = 1000;
 const getId = () => `${id++}`;
 
@@ -68,6 +77,10 @@ const Builder: React.FC<BuilderProps> = ({ onBackToHome }) => {
   const [isApiKeyModalOpen, setIsApiKeyModalOpen] = useState(false);
   const [isConnected, setIsConnected] = useState(false);
 
+  // Model Selection State
+  const [selectedModel, setSelectedModel] = useState(DEFAULT_MODEL);
+  const [isModelMenuOpen, setIsModelMenuOpen] = useState(false);
+
   // Activity Log State
   const [logs, setLogs] = useState<LogEntry[]>([]);
 
@@ -207,9 +220,9 @@ const Builder: React.FC<BuilderProps> = ({ onBackToHome }) => {
     if (!isConnected) { setIsApiKeyModalOpen(true); return; }
     setIsGeneratingPrompt(true);
     setIsPromptViewerOpen(true);
-    addLog("Generating PyTorch code prompt...", 'info');
+    addLog(`Generating PyTorch code prompt using ${AVAILABLE_MODELS.find(m => m.id === selectedModel)?.name}...`, 'info');
     try {
-        const promptText = await generateRefinedPrompt(nodes, edges);
+        const promptText = await generateRefinedPrompt(nodes, edges, selectedModel);
         setGeneratedPrompt(promptText);
         addLog("Code prompt generated successfully.", 'success');
     } catch (e) {
@@ -223,9 +236,9 @@ const Builder: React.FC<BuilderProps> = ({ onBackToHome }) => {
   const handleValidate = async () => {
     if (!isConnected) { setIsApiKeyModalOpen(true); return; }
     setValidationMsg("Validating...");
-    addLog("Running architecture validation...", 'info');
+    addLog(`Running architecture validation using ${AVAILABLE_MODELS.find(m => m.id === selectedModel)?.name}...`, 'info');
     try {
-        const result = await validateArchitecture(nodes, edges);
+        const result = await validateArchitecture(nodes, edges, selectedModel);
         setValidationMsg(result);
         if (result.toLowerCase().includes("valid") && !result.toLowerCase().includes("invalid") && result.length < 50) {
             setTimeout(() => setValidationMsg(null), 5000);
@@ -244,7 +257,7 @@ const Builder: React.FC<BuilderProps> = ({ onBackToHome }) => {
     setIsSuggestionsOpen(true);
     setIsSuggestionsLoading(true);
     try {
-      const suggestions = await getArchitectureSuggestions(nodes, edges);
+      const suggestions = await getArchitectureSuggestions(nodes, edges, selectedModel);
       setSuggestionsText(suggestions);
     } catch (error) {
       setSuggestionsText("Failed to get suggestions. API Error.");
@@ -257,9 +270,9 @@ const Builder: React.FC<BuilderProps> = ({ onBackToHome }) => {
   const handleImplementSuggestions = async () => {
     if (!suggestionsText || isImplementingSuggestions) return;
     setIsImplementingSuggestions(true);
-    addLog("AI Agents implementing suggestions...", 'info');
+    addLog(`AI Agents implementing suggestions using ${AVAILABLE_MODELS.find(m => m.id === selectedModel)?.name}...`, 'info');
     try {
-      const result = await implementArchitectureSuggestions(nodes, edges, suggestionsText);
+      const result = await implementArchitectureSuggestions(nodes, edges, suggestionsText, selectedModel);
       if (result && result.nodes) {
          // Fix types and defaults safely
          const processedNodes = result.nodes.map((n: any) => {
@@ -329,8 +342,9 @@ const Builder: React.FC<BuilderProps> = ({ onBackToHome }) => {
         }
       }));
 
-      const hydratedEdges = clonedEdges.map((e: any) => ({
+      const hydratedEdges = clonedEdges.map((e: any, idx: number) => ({
         ...e,
+        id: e.id || `e-template-${idx}-${Math.random().toString(36).substr(2, 4)}`,
         animated: true,
         style: { stroke: '#94a3b8' }
       }));
@@ -412,6 +426,44 @@ const Builder: React.FC<BuilderProps> = ({ onBackToHome }) => {
               <Key size={16} className={isConnected ? "text-emerald-500" : "text-slate-400"} />
               <span className="hidden md:inline">{isConnected ? 'Key Active' : 'Add Key'}</span>
             </button>
+
+            {/* Model Selection Button */}
+            <div className="relative">
+              <button 
+                onClick={() => setIsModelMenuOpen(!isModelMenuOpen)}
+                className="flex items-center gap-2 px-3 py-2 bg-slate-800/80 backdrop-blur hover:bg-slate-700 border border-slate-700 rounded-lg text-slate-200 text-xs md:text-sm font-medium transition-colors shadow-lg"
+                title="Select Gemini Model"
+              >
+                <Cpu size={16} className="text-blue-400" />
+                <span className="hidden md:inline">{AVAILABLE_MODELS.find(m => m.id === selectedModel)?.name.split(' ')[1] || 'Model'}</span>
+                <ChevronDown size={14} className={`text-slate-500 transition-transform ${isModelMenuOpen ? 'rotate-180' : ''}`} />
+              </button>
+              
+              {isModelMenuOpen && (
+                <div className="absolute right-0 mt-2 w-64 bg-slate-900 border border-slate-700 rounded-xl shadow-2xl z-[100] overflow-hidden animate-in fade-in slide-in-from-top-2 duration-200">
+                  <div className="p-2">
+                    {AVAILABLE_MODELS.map((model) => (
+                      <button
+                        key={model.id}
+                        onClick={() => {
+                          setSelectedModel(model.id);
+                          setIsModelMenuOpen(false);
+                          addLog(`Switched to model: ${model.name}`, 'info');
+                        }}
+                        className={`w-full text-left p-3 rounded-lg transition-colors ${selectedModel === model.id ? 'bg-blue-600/20 border border-blue-500/30' : 'hover:bg-slate-800 border border-transparent'}`}
+                      >
+                        <div className="flex items-center justify-between">
+                          <span className={`font-semibold text-sm ${selectedModel === model.id ? 'text-blue-400' : 'text-slate-200'}`}>{model.name}</span>
+                          {selectedModel === model.id && <Check size={14} className="text-blue-400" />}
+                        </div>
+                        <div className="text-xs text-slate-500 mt-0.5">{model.description}</div>
+                      </button>
+                    ))}
+                  </div>
+                </div>
+              )}
+            </div>
+
             <button 
               onClick={handleGetSuggestions}
               className="flex items-center gap-2 px-3 py-2 bg-slate-800/80 backdrop-blur hover:bg-slate-700 border border-slate-700 rounded-lg text-slate-200 text-xs md:text-sm font-medium transition-colors shadow-lg"
@@ -537,6 +589,7 @@ const Builder: React.FC<BuilderProps> = ({ onBackToHome }) => {
         onClose={() => setIsAIBuilderOpen(false)}
         onApply={handleApplyAIBuild}
         currentNodes={nodes}
+        selectedModel={selectedModel}
       />
 
       <FixerModal
@@ -546,6 +599,7 @@ const Builder: React.FC<BuilderProps> = ({ onBackToHome }) => {
         errorMsg={validationMsg || ''}
         nodes={nodes}
         edges={edges}
+        selectedModel={selectedModel}
       />
 
       <ApiKeyModal