Update services/geminiService.ts

services/geminiService.ts

@@ -1,8 +1,67 @@
-import { GoogleGenAI } from "@google/genai";
+import { GoogleGenAI, Type } from "@google/genai";
 import { Node, Edge } from 'reactflow';
 import { NodeData, LayerType } from '../types';
 import { LAYER_DEFINITIONS } from '../constants';
 
+// --- SCHEMAS ---
+const graphResponseSchema = {
+  type: Type.OBJECT,
+  properties: {
+    reasoning: { 
+        type: Type.STRING, 
+        description: "Brief explanation of architectural choices or fixes applied." 
+    },
+    nodes: {
+      type: Type.ARRAY,
+      items: {
+        type: Type.OBJECT,
+        properties: {
+          id: { type: Type.STRING },
+          type: { type: Type.STRING, description: "Must be 'custom' for ReactFlow compatibility." },
+          position: {
+            type: Type.OBJECT,
+            properties: {
+              x: { type: Type.NUMBER },
+              y: { type: Type.NUMBER }
+            },
+            required: ["x", "y"]
+          },
+          data: {
+            type: Type.OBJECT,
+            properties: {
+              type: { type: Type.STRING, description: "The LayerType enum value." },
+              label: { type: Type.STRING },
+              params: { type: Type.OBJECT, description: "Hyperparameters for the layer." }
+            },
+            required: ["type", "label"]
+          }
+        },
+        required: ["id", "type", "position", "data"]
+      }
+    },
+    edges: {
+      type: Type.ARRAY,
+      items: {
+        type: Type.OBJECT,
+        properties: {
+          id: { type: Type.STRING },
+          source: { type: Type.STRING },
+          target: { type: Type.STRING }
+        },
+        required: ["source", "target"]
+      }
+    }
+  },
+  required: ["nodes", "edges"]
+};
+
+/**
+ * Helper: Cleans a JSON string that might be wrapped in markdown backticks.
+ */
+const cleanJsonString = (str: string): string => {
+    return str.replace(/```json/g, '').replace(/```/g, '').trim();
+};
+
 // Key Management
 const getEnvKey = () => {
   // In Development: Returns key from .env file if available
@@ -33,7 +92,7 @@ const getAiClient = () => {
   return new GoogleGenAI({ apiKey: key });
 };
 
-const MODEL_NAME = 'gemini-2.5-flash';
+export const DEFAULT_MODEL = 'gemini-2.5-flash';
 
 export type AgentStatus = 'idle' | 'architect' | 'critic' | 'refiner' | 'debugger' | 'patcher' | 'complete' | 'error';
 
@@ -100,46 +159,44 @@ const buildRawPrompt = (nodes: Node<NodeData>[], edges: Edge[]): string => {
  * Generates a polished, professional prompt using the AI.
  * It takes the raw hardcoded spec and asks the AI to format it perfectly for a coding LLM.
  */
-export const generateRefinedPrompt = async (nodes: Node<NodeData>[], edges: Edge[]): Promise<string> => {
+export const generateRefinedPrompt = async (nodes: Node<NodeData>[], edges: Edge[], model: string = DEFAULT_MODEL): Promise<string> => {
     const ai = getAiClient();
     const rawSpec = buildRawPrompt(nodes, edges);
 
-    const systemPrompt = `
-    You are an expert AI Prompt Engineer for Deep Learning.
-    Your goal is to take a raw, technical neural network specification and rewrite it into a
-    perfect, professional, and detailed prompt that another AI (like a coding assistant) could use to write flawless PyTorch code.
-
-    Input Raw Specification:
-    ${rawSpec}
-
-    Instructions:
-    1. Start the output with: "You are an expert Deep Learning Engineer. Please write a complete, runnable PyTorch model code for the following neural network architecture:"
-    2. Create a section "Architecture Specification".
-    3. List "1. Layers (Nodes)" cleanly. Include ID, Type, and Parameters.
-    4. List "2. Connectivity (Forward Pass Flow)" cleanly. 
-       - Explicitly describe merge points (e.g. "Node X receives inputs from A and B. Handle this merge...").
-       - Note specific handling for complex layers like CrossAttention (needs Query + Key/Value) or SAM Decoders.
-    5. Create a section "Implementation Requirements" with standard PyTorch best practices (nn.Module, forward method, correct shapes).
-    6. If CUSTOM_CODE_DEFINITION or CUSTOM_IMPORTS are present, explicitly instruct the coder to include them verbatim or use them as reference.
-    7. Do NOT write the Python code yourself. Write the PROMPT that asks for the code.
-    8. Ensure the tone is technical and precise.
-
-    Return ONLY the generated prompt text.
-    `;
-
     try {
         const response = await ai.models.generateContent({
-            model: MODEL_NAME,
-            contents: systemPrompt,
+            model: model,
+            contents: `Input Raw Specification:\n${rawSpec}`,
+            config: {
+                systemInstruction: `You are an expert AI Prompt Engineer specializing in Deep Learning Architecture.
+Your goal is to transform a raw graph specification into a high-fidelity, professional prompt for a coding LLM.
+You have access to Google Search to look up the latest PyTorch best practices or SOTA implementation details to include in the prompt.
+
+Instructions:
+1. Start with: "You are a world-class Deep Learning Engineer. Implement the following PyTorch model with precision:"
+2. Section "Architecture Details": List all layers with their IDs, Types, and specific hyperparameters.
+3. Section "Data Flow & Connectivity": Describe the forward pass step-by-step. 
+   - Explicitly mention skip connections, residual additions, and concatenation points.
+   - For multi-input layers (like CrossAttention or Add), specify exactly which node IDs provide the inputs.
+4. Section "Implementation Requirements":
+   - Use idiomatic PyTorch (nn.Module).
+   - Ensure input/output shapes are documented in comments.
+   - Include proper initialization (e.g., Kaiming or Xavier) where appropriate.
+   - Handle potential shape mismatches (e.g., adding a Flatten layer before Linear if needed).
+5. If custom code is provided in the spec, ensure it is integrated correctly.
+6. Do NOT write the code yourself. Write the PROMPT that will guide another AI to write the code.
+7. Maintain a highly technical, rigorous, and clear tone.`,
+                tools: [{ googleSearch: {} }]
+            }
         });
         return response.text.trim();
     } catch (error) {
         console.error("Prompt refinement failed:", error);
-        throw error; // Re-throw so caller can handle auth errors
+        throw error;
     }
 };
 
-export const validateArchitecture = async (nodes: Node<NodeData>[], edges: Edge[]): Promise<string> => {
+export const validateArchitecture = async (nodes: Node<NodeData>[], edges: Edge[], model: string = DEFAULT_MODEL): Promise<string> => {
    const ai = getAiClient();
    const graphRepresentation = {
     nodes: nodes.map(n => ({
@@ -153,24 +210,24 @@ export const validateArchitecture = async (nodes: Node<NodeData>[], edges: Edge[
     }))
   };
 
-  const prompt = `
-  Analyze this neural network architecture graph for validity.
-  Graph: ${JSON.stringify(graphRepresentation)}
-  
-  Check for:
-  1. Shape mismatches (e.g., Conv2D output to Linear without Flatten).
-  2. Disconnected components.
-  3. Logical errors (e.g., MaxPool after Output).
-  4. Merge layer correctness (Concat/Add needs multiple inputs).
-  5. GenAI correctness (e.g., CrossAttention needs 2 inputs, VLM projection dims match).
-  
-  Return a concise report. If valid, say "Architecture is valid.". If invalid, list specific errors and suggest fixes.
-  `;
-
   try {
     const response = await ai.models.generateContent({
-      model: MODEL_NAME,
-      contents: prompt,
+      model: model,
+      contents: `Graph to Analyze: ${JSON.stringify(graphRepresentation)}`,
+      config: {
+          systemInstruction: `You are a Senior Deep Learning Validator. Your role is to find structural and logical flaws in neural network architectures.
+
+Checklist:
+1. Dimensional Consistency: Do Conv2D/Conv3D layers have appropriate pooling or flattening before Linear layers?
+2. Connectivity: Are there any orphaned nodes? Is there a path from Input to Output?
+3. Layer Logic: Are activations placed correctly? (e.g., no ReLU after a Softmax).
+4. Merge Operations: Do Add/Concat layers have at least 2 inputs? Do they have compatible shapes?
+5. GenAI Patterns: Does Attention have Query/Key/Value paths? Do VLM projections match the LLM backbone dimensions?
+
+Output Format:
+- If valid: "Architecture is valid."
+- If invalid: A bulleted list of "CRITICAL ERRORS" and "SUGGESTED FIXES".`
+      }
     });
     return response.text;
   } catch (error) {
@@ -181,7 +238,7 @@ export const validateArchitecture = async (nodes: Node<NodeData>[], edges: Edge[
 /**
  * Gets AI suggestions for improving the architecture.
  */
-export const getArchitectureSuggestions = async (nodes: Node<NodeData>[], edges: Edge[]): Promise<string> => {
+export const getArchitectureSuggestions = async (nodes: Node<NodeData>[], edges: Edge[], model: string = DEFAULT_MODEL): Promise<string> => {
    const ai = getAiClient();
    const graphRepresentation = {
     nodes: nodes.map(n => ({
@@ -196,26 +253,24 @@ export const getArchitectureSuggestions = async (nodes: Node<NodeData>[], edges:
     }))
   };
 
-  const prompt = `
-  You are a Senior Deep Learning Architect. Review the following neural network architecture graph.
-  
-  Graph Structure:
-  ${JSON.stringify(graphRepresentation, null, 2)}
-  
-  Task: Provide 3 to 5 concrete, actionable suggestions to improve this model. 
-  Focus on:
-  - Modern best practices (e.g., using LayerNorm vs BatchNorm in Transformers, SwiGLU vs ReLU).
-  - Architecture efficiency and parameter count optimization.
-  - Potential bottlenecks or vanishing gradient risks.
-  - Adding residuals or skip connections if the model is deep.
-  
-  Format the output as a clean bulleted list. Keep it concise, professional and helpful.
-  `;
-
   try {
     const response = await ai.models.generateContent({
-      model: MODEL_NAME,
-      contents: prompt,
+      model: model,
+      contents: `Current Graph Architecture: ${JSON.stringify(graphRepresentation, null, 2)}`,
+      config: {
+          systemInstruction: `You are a Senior Deep Learning Architect. Your goal is to provide elite-level optimization suggestions for neural network architectures.
+You have access to Google Search to research the latest SOTA (State of the Art) components and research papers.
+
+Focus Areas:
+1. Efficiency: Reduce parameter count or computational complexity without sacrificing accuracy.
+2. Modernity: Suggest state-of-the-art components (e.g., FlashAttention, SwiGLU, RMSNorm).
+3. Robustness: Identify risks of vanishing/exploding gradients and suggest residuals or normalization.
+4. Scalability: Suggest ways to make the model more modular or scalable.
+
+Format:
+Provide 3-5 concise, bulleted suggestions. Each suggestion should include a "Why" (the technical benefit).`,
+          tools: [{ googleSearch: {} }]
+      }
     });
     return response.text;
   } catch (error) {
@@ -229,7 +284,8 @@ export const getArchitectureSuggestions = async (nodes: Node<NodeData>[], edges:
 export const implementArchitectureSuggestions = async (
   nodes: Node<NodeData>[], 
   edges: Edge[], 
-  suggestions: string
+  suggestions: string,
+  model: string = DEFAULT_MODEL
 ): Promise<{ nodes: any[], edges: any[] }> => {
   const ai = getAiClient();
   
@@ -247,29 +303,23 @@ export const implementArchitectureSuggestions = async (
     }))
   };
 
-  const prompt = `
-  You are a Senior Implementation Engineer.
-  Task: Apply the following architectural suggestions to the provided graph JSON.
-  
-  Current Graph:
-  ${JSON.stringify(graphRepresentation)}
-  
-  Suggestions to Implement:
-  "${suggestions}"
-  
-  Instructions:
-  1. Modify the nodes and edges to incorporate the suggestions.
-  2. Maintain the layout (x, y positions) as best as possible, offsetting new nodes if added.
-  3. Ensure all LayerTypes are valid from the standard schema.
-  4. Return the complete, updated JSON with "nodes" and "edges" arrays.
-  5. Return ONLY raw JSON.
-  `;
-
   try {
     const response = await ai.models.generateContent({
-      model: MODEL_NAME,
-      contents: prompt,
-      config: { responseMimeType: "application/json" }
+        model: model,
+        contents: `Current Graph Architecture: ${JSON.stringify(graphRepresentation)}\nSuggestions to Implement: "${suggestions}"`,
+        config: { 
+            systemInstruction: `You are a Senior Implementation Engineer specializing in Deep Learning. 
+Your task is to modify the provided neural network graph JSON to incorporate specific architectural improvements.
+
+Rules:
+1. Maintain the relative layout (x, y positions). If adding nodes, place them logically between existing ones.
+2. Use ONLY valid LayerTypes from the schema.
+3. Ensure all new nodes have unique IDs.
+4. Ensure the resulting graph is fully connected and logically sound.
+5. Preserve existing node parameters unless the suggestion explicitly requires changing them.`,
+            responseMimeType: "application/json",
+            responseSchema: graphResponseSchema
+        }
     });
     const result = JSON.parse(response.text.trim());
     return sanitizeGraph(result);
@@ -284,47 +334,72 @@ export const implementArchitectureSuggestions = async (
  * Checks if 'type' is valid. If not, converts it to a CUSTOM layer to prevent crashes.
  */
 const sanitizeGraph = (graphJson: { nodes: any[], edges: any[] }) => {
-    if (!graphJson || !graphJson.nodes) return graphJson;
+    if (!graphJson) return graphJson;
     
-    graphJson.nodes = graphJson.nodes.map(node => {
-        // AI might return type in data.type or top-level type
-        let rawType = node.data?.type || node.type || 'Identity';
-        
-        // Ensure rawType is a string
-        if (typeof rawType !== 'string') rawType = 'Identity';
-
-        // Check if this type exists in our known definitions
-        const isValid = Object.values(LayerType).includes(rawType as LayerType) && LAYER_DEFINITIONS[rawType as LayerType];
-
-        if (!isValid) {
-            console.warn(`Sanitizing unknown layer type: ${rawType}. Converting to CustomLayer.`);
+    if (graphJson.nodes) {
+        graphJson.nodes = graphJson.nodes.map(node => {
+            // AI might return type in data.type or top-level type
+            let rawType = node.data?.type || node.type || 'Identity';
+            
+            // Ensure rawType is a string
+            if (typeof rawType !== 'string') rawType = 'Identity';
+
+            // Check if this type exists in our known definitions
+            const isValid = Object.values(LayerType).includes(rawType as LayerType) && LAYER_DEFINITIONS[rawType as LayerType];
+
+            if (!isValid) {
+                console.warn(`Sanitizing unknown layer type: ${rawType}. Converting to CustomLayer.`);
+                return {
+                    ...node,
+                    type: 'custom', // ReactFlow type
+                    data: {
+                        ...node.data,
+                        type: LayerType.CUSTOM,
+                        label: node.data?.label || rawType,
+                        params: {
+                            ...(node.data?.params || {}),
+                            class_name: rawType, // Store original type name here
+                            args: JSON.stringify(node.data?.params || {}).slice(0, 100) // Rough preserve of args
+                        }
+                    }
+                };
+            }
+            
+            // Ensure standard structure for valid nodes
             return {
                 ...node,
-                type: 'custom', // ReactFlow type
+                type: 'custom',
                 data: {
                     ...node.data,
-                    type: LayerType.CUSTOM,
-                    label: node.data?.label || rawType,
-                    params: {
-                        ...(node.data?.params || {}),
-                        class_name: rawType, // Store original type name here
-                        args: JSON.stringify(node.data?.params || {}).slice(0, 100) // Rough preserve of args
-                    }
+                    type: rawType,
+                    params: node.data?.params || {}
                 }
             };
-        }
-        
-        // Ensure standard structure for valid nodes
-        return {
-            ...node,
-            type: 'custom',
-            data: {
-                ...node.data,
-                type: rawType,
-                params: node.data?.params || {}
+        });
+    }
+
+    if (graphJson.edges) {
+        const seenIds = new Set<string>();
+        graphJson.edges = graphJson.edges.map((edge, idx) => {
+            // AI might return edge without ID or with duplicate ID
+            let baseId = edge.id || `e-${edge.source}-${edge.target}`;
+            
+            // Ensure uniqueness within this graph
+            let uniqueId = baseId;
+            let counter = 1;
+            while (seenIds.has(uniqueId)) {
+                uniqueId = `${baseId}-${idx}-${counter++}`;
             }
-        };
-    });
+            seenIds.add(uniqueId);
+
+            return {
+                ...edge,
+                id: uniqueId,
+                animated: true,
+                style: { stroke: '#94a3b8' }
+            };
+        });
+    }
 
     return graphJson;
 };
@@ -339,7 +414,8 @@ const sanitizeGraph = (graphJson: { nodes: any[], edges: any[] }) => {
 export const generateGraphWithAgents = async (
   userPrompt: string, 
   currentNodes: Node<NodeData>[] = [],
-  onStatusUpdate: (status: AgentStatus, log: string) => void
+  onStatusUpdate: (status: AgentStatus, log: string) => void,
+  model: string = DEFAULT_MODEL
 ): Promise<{ nodes: any[], edges: any[] } | null> => {
   const ai = getAiClient();
 
@@ -356,34 +432,34 @@ export const generateGraphWithAgents = async (
     ? `Current Graph Context: ${JSON.stringify(currentNodes.map(n => ({ id: n.id, type: n.data.type, label: n.data.label })))}`
     : "Starting from scratch.";
 
-  const architectPrompt = `
-    Role: Senior Neural Network Architect.
-    Task: Create a preliminary graph layout (JSON) for the user request.
-    User Request: "${userPrompt}"
-    Context: ${context}
-    
-    Available Layers: ${Object.keys(LAYER_DEFINITIONS).join(', ')}
-    Schema Reference: ${schemaStr}
-    
-    Requirements:
-    1. Output valid JSON with "nodes" and "edges" arrays.
-    2. "nodes": { id, type='custom', position:{x,y}, data:{ type: LayerType, label: string, params: {} } }
-    3. Use correct LayerTypes from enum.
-    4. Layout nodes vertically (y+150 each step).
-    5. Connect edges logically.
-    6. If multi-input/output, arrange horizontally.
-    
-    Return ONLY raw JSON.
-  `;
+  const architectSystemInstruction = `You are a Senior Neural Network Architect. 
+Your task is to design a high-quality neural network graph based on user requirements.
+
+Design Principles:
+1. Hierarchical Flow: Arrange nodes from top (Input) to bottom (Output) with y-offsets of +150.
+2. Modern Patterns: Use residuals/skip connections for deep networks. Prefer LayerNorm/RMSNorm for Transformers.
+3. Spatial Clarity: If a layer has multiple branches, spread them horizontally (x-offsets).
+4. Schema Compliance: Use ONLY valid LayerTypes and parameter names provided in the schema.
+
+Output Format:
+Return a JSON object with:
+- "reasoning": A brief string explaining your architectural choices.
+- "nodes": Array of { id, type='custom', position:{x,y}, data:{ type, label, params:{} } }
+- "edges": Array of { id, source, target }`;
 
   let draftJsonStr = "";
+  let architectResponse;
   try {
-    const response = await ai.models.generateContent({
-        model: MODEL_NAME,
-        contents: architectPrompt,
-        config: { responseMimeType: "application/json" }
+    architectResponse = await ai.models.generateContent({
+        model: model,
+        contents: `User Request: "${userPrompt}"\nContext: ${context}\nSchema: ${schemaStr}`,
+        config: { 
+            systemInstruction: architectSystemInstruction,
+            // REMOVED responseMimeType and responseSchema as they conflict with tools (googleSearch/urlContext)
+            tools: [{ googleSearch: {} }, { urlContext: {} }]
+        }
     });
-    draftJsonStr = response.text.trim();
+    draftJsonStr = cleanJsonString(architectResponse.text);
   } catch (e) {
       throw e;
   }
@@ -391,29 +467,35 @@ export const generateGraphWithAgents = async (
   // --- Step 2: Critic ---
   onStatusUpdate('critic', 'Critic is reviewing architecture for flaws...');
   
-  const criticPrompt = `
-    Role: Senior Lead Reviewer.
-    Task: Critique the following neural network architecture draft.
-    User Request: "${userPrompt}"
-    Draft Architecture: ${draftJsonStr}
-    
-    Check for:
-    - Shape mismatches (e.g. 3D output into 2D input without flattening)
-    - Logical connection errors
-    - Missing essential layers (e.g. Activations, Normalization)
-    - Parameter errors (e.g. kernel size too large)
-    - Compliance with user request
-    
-    Output a concise paragraph describing specific improvements needed. If perfect, say "No changes needed".
-  `;
+  const criticSystemInstruction = `You are a Lead Architecture Reviewer. 
+Critique the provided neural network draft for technical soundness, efficiency, and adherence to the user's request.
+You have access to Google Search to verify SOTA (State of the Art) practices if needed.
+
+Check for:
+- Missing critical components (e.g., no activation after Conv, no pooling before FC).
+- Bottlenecks (e.g., massive dimensionality jumps).
+- Redundancy or inefficient paths.
+- Logical errors in connectivity.
+
+If the architecture is excellent, say "No changes needed". Otherwise, provide a concise list of required improvements.`;
 
   let critique = "";
+  let criticResponse;
   try {
-     const response = await ai.models.generateContent({
-        model: MODEL_NAME,
-        contents: criticPrompt,
+     criticResponse = await ai.models.generateContent({
+        model: model,
+        // Shared Context: Pass the architect's response as a previous turn to be token efficient
+        contents: [
+            { role: 'user', parts: [{ text: `User Request: "${userPrompt}"\nSchema: ${schemaStr}` }] },
+            { role: 'model', parts: [{ text: draftJsonStr }] },
+            { role: 'user', parts: [{ text: "Please critique this architecture." }] }
+        ],
+        config: { 
+            systemInstruction: criticSystemInstruction,
+            tools: [{ googleSearch: {} }]
+        }
     });
-    critique = response.text.trim();
+    critique = criticResponse.text.trim();
   } catch (e) {
       console.warn("Critic agent failed, proceeding with draft.");
       critique = "No critique available.";
@@ -422,33 +504,35 @@ export const generateGraphWithAgents = async (
   // --- Step 3: Refiner ---
   onStatusUpdate('refiner', 'Refiner is applying fixes and finalizing...');
   
-  const refinerPrompt = `
-    Role: Lead Engineer.
-    Task: Finalize the JSON architecture based on the critique.
-    
-    Draft: ${draftJsonStr}
-    Critique: "${critique}"
-    
-    Instructions:
-    1. Apply the fixes suggested in the critique.
-    2. Ensure the JSON structure is strictly { "nodes": [...], "edges": [...] }.
-    3. Ensure all node IDs are unique strings.
-    4. Ensure parameter values match the schema types.
-    5. Ensure "type" in top level node object is always 'custom'.
-    
-    Return ONLY the final JSON.
-  `;
+  const refinerSystemInstruction = `You are a Lead Implementation Engineer. 
+Finalize the neural network JSON by applying the Critic's feedback to the Architect's draft.
+
+Requirements:
+1. Strictly follow the JSON schema: { "nodes": [...], "edges": [...] }.
+2. Ensure all node IDs are unique.
+3. Ensure all LayerTypes match the enum exactly.
+4. Fix any shape or logic errors identified in the critique.
+5. Maintain a clean, readable visual layout.
+
+Return ONLY the final JSON object.`;
 
   try {
     const response = await ai.models.generateContent({
-        model: MODEL_NAME,
-        contents: refinerPrompt,
-        config: { responseMimeType: "application/json" }
+        model: model,
+        // Shared Context: Pass the entire history for maximum consistency
+        contents: [
+            { role: 'user', parts: [{ text: `User Request: "${userPrompt}"\nSchema: ${schemaStr}` }] },
+            { role: 'model', parts: [{ text: draftJsonStr }] },
+            { role: 'user', parts: [{ text: `Critique: "${critique}"\n\nApply the critique and provide the final JSON.` }] }
+        ],
+        config: { 
+            systemInstruction: refinerSystemInstruction,
+            // REMOVED responseMimeType and responseSchema as they conflict with previous tool use in history
+        }
     });
-    const finalJson = JSON.parse(response.text.trim());
+    const finalJson = JSON.parse(cleanJsonString(response.text));
     onStatusUpdate('complete', 'Architecture built successfully!');
     
-    // SANITIZE: Prevent UI crashes by handling hallucinated types
     return sanitizeGraph(finalJson);
   } catch (e) {
      throw new Error("Refiner agent failed to parse final JSON.");
@@ -461,80 +545,100 @@ export const generateGraphWithAgents = async (
  */
 export const generateCodeWithAgents = async (
   promptText: string,
-  onStatusUpdate: (status: AgentStatus, log: string) => void
+  onStatusUpdate: (status: AgentStatus, log: string) => void,
+  model: string = DEFAULT_MODEL
 ): Promise<string> => {
   const ai = getAiClient();
 
-  // --- Step 1: Coder (Architect) ---
+  // --- Step 1: Coder ---
   onStatusUpdate('architect', 'Coder Agent is writing initial PyTorch implementation...');
 
-  const coderPrompt = `
-    Role: Senior Deep Learning Engineer.
-    Task: Write complete, runnable PyTorch code based on the following architecture prompt.
-    Prompt: "${promptText}"
-    
-    Requirements:
-    - Use torch.nn.Module
-    - Include all necessary imports
-    - Handle forward pass logic exactly as described
-    - Include a 'if __name__ == "__main__":' block to test with dummy data
-    
-    Return ONLY the Python code. No markdown formatting.
-  `;
+  const coderSystemInstruction = `You are a Senior Deep Learning Engineer. 
+Your task is to write clean, modular, and production-ready PyTorch code.
+
+Coding Standards:
+1. Use nn.Module for the main model.
+2. Include docstrings for the class and forward method.
+3. Add comments indicating the expected tensor shapes at each major step.
+4. Use descriptive variable names.
+5. Include a robust 'if __name__ == "__main__":' block that instantiates the model and runs a dummy forward pass with torch.randn.
+6. Handle custom layer logic or external imports if specified.
+
+Return ONLY the Python code. No markdown backticks.`;
 
   let draftCode = "";
   try {
       const response = await ai.models.generateContent({
-        model: MODEL_NAME,
-        contents: coderPrompt
+        model: model,
+        contents: `Request: "${promptText}"`,
+        config: { 
+            systemInstruction: coderSystemInstruction,
+            tools: [{ googleSearch: {} }, { urlContext: {} }]
+        }
       });
       draftCode = response.text.trim().replace(/```python/g, '').replace(/```/g, '');
   } catch(e) {
       throw e;
   }
 
-  // --- Step 2: Reviewer (Critic) ---
+  // --- Step 2: Reviewer ---
   onStatusUpdate('critic', 'Reviewer Agent is analyzing code for bugs and optimization...');
-  const reviewPrompt = `
-    Role: Code Reviewer.
-    Task: Review the following PyTorch code for errors, shape mismatches, or style issues.
-    Code:
-    ${draftCode}
-    
-    Original Prompt Request: "${promptText}"
-    
-    Output a concise critique. If perfect, say "No changes needed".
-  `;
-  
+
+  const reviewerSystemInstruction = `You are a Senior Code Reviewer. 
+Analyze the provided PyTorch code for:
+- Syntax errors or missing imports.
+- Logical bugs (e.g., wrong dimension in cat/stack).
+- Missing super().__init__() calls.
+- Inefficient implementations.
+- Adherence to the original architecture request.
+You have access to Google Search to verify PyTorch documentation or SOTA implementation details.
+
+Provide a concise, technical critique. If the code is perfect, say "No changes needed".`;
+
   let critique = "";
   try {
       const response = await ai.models.generateContent({
-        model: MODEL_NAME,
-        contents: reviewPrompt
+        model: model,
+        contents: [
+            { role: 'user', parts: [{ text: `Original Request: "${promptText}"` }] },
+            { role: 'model', parts: [{ text: draftCode }] },
+            { role: 'user', parts: [{ text: "Please review this code." }] }
+        ],
+        config: { 
+            systemInstruction: reviewerSystemInstruction,
+            tools: [{ googleSearch: {} }]
+        }
       });
       critique = response.text.trim();
   } catch(e) {
       critique = "No critique available.";
   }
 
-  // --- Step 3: Polisher (Refiner) ---
+  // --- Step 3: Polisher ---
   onStatusUpdate('refiner', 'Polisher Agent is finalizing the codebase...');
-  const polisherPrompt = `
-    Role: Senior Software Engineer.
-    Task: Refine the PyTorch code based on the critique.
-    
-    Draft Code: ${draftCode}
-    Critique: ${critique}
-    
-    Return ONLY the final Python code. No markdown formatting.
-  `;
+
+  const polisherSystemInstruction = `You are a Principal Software Engineer. 
+Refine the PyTorch code by incorporating the Reviewer's feedback.
+
+Goals:
+1. Fix all bugs and style issues.
+2. Ensure the code is strictly valid Python.
+3. Maintain all docstrings and shape comments.
+4. Ensure the test block works perfectly.
+
+Return ONLY the final Python code. No markdown backticks.`;
 
   try {
       const response = await ai.models.generateContent({
-        model: MODEL_NAME,
-        contents: polisherPrompt
+        model: model,
+        contents: [
+            { role: 'user', parts: [{ text: `Original Request: "${promptText}"` }] },
+            { role: 'model', parts: [{ text: draftCode }] },
+            { role: 'user', parts: [{ text: `Critique: ${critique}\n\nApply the critique and provide the final polished code.` }] }
+        ],
+        config: { systemInstruction: polisherSystemInstruction }
       });
-      let finalCode = response.text.trim().replace(/```python/g, '').replace(/```/g, '');
+      let finalCode = cleanJsonString(response.text);
       onStatusUpdate('complete', 'Code generation complete!');
       return finalCode;
   } catch(e) {
@@ -549,7 +653,8 @@ export const fixArchitectureErrors = async (
   nodes: Node<NodeData>[],
   edges: Edge[],
   errorMsg: string,
-  onStatusUpdate: (status: AgentStatus, log: string) => void
+  onStatusUpdate: (status: AgentStatus, log: string) => void,
+  model: string = DEFAULT_MODEL
 ): Promise<{ nodes: any[], edges: any[] } | null> => {
   const ai = getAiClient();
   const graphJson = JSON.stringify({
@@ -560,21 +665,24 @@ export const fixArchitectureErrors = async (
   // --- Step 1: Debugger ---
   onStatusUpdate('debugger', 'Debugger Agent is analyzing the error trace...');
   
-  const debuggerPrompt = `
-    Role: Senior Systems Debugger.
-    Task: Analyze the architecture graph and the reported error to pinpoint the root cause.
-    
-    Graph: ${graphJson}
-    Error Message: "${errorMsg}"
-    
-    Output a technical analysis of exactly what is wrong (e.g. "Node A connects to Node B but shapes [X] and [Y] are incompatible").
-  `;
-
   let debugAnalysis = "";
   try {
     const response = await ai.models.generateContent({
-      model: MODEL_NAME,
-      contents: debuggerPrompt,
+      model: model,
+      contents: `Graph: ${graphJson}\nError Message: "${errorMsg}"`,
+      config: {
+          systemInstruction: `You are a Senior Systems Debugger specializing in Deep Learning. 
+Analyze the provided graph and error message to identify the root cause of the failure.
+You have access to Google Search to look up specific PyTorch error traces or layer compatibility issues.
+
+Focus on:
+- Shape mismatches between connected layers.
+- Incorrect parameter values.
+- Missing required layers (e.g., Flatten before Linear).
+
+Output a concise, technical explanation of the bug.`,
+          tools: [{ googleSearch: {} }]
+      }
     });
     debugAnalysis = response.text.trim();
   } catch (e) {
@@ -584,57 +692,77 @@ export const fixArchitectureErrors = async (
   // --- Step 2: Architect ---
   onStatusUpdate('architect', 'Architect Agent is planning the fix...');
   
-  const architectPrompt = `
-    Role: Solution Architect.
-    Task: Propose a specific fix for the identified issue.
-    
-    Issue Analysis: ${debugAnalysis}
-    
-    Instructions:
-    - Determine if nodes need to be added (e.g. Flatten, Reshape), removed, or reconnected.
-    - Determine if parameters need changing.
-    
-    Output the plan in clear steps.
-  `;
-
   let fixPlan = "";
   try {
     const response = await ai.models.generateContent({
-      model: MODEL_NAME,
-      contents: architectPrompt,
+      model: model,
+      contents: `Issue Analysis: ${debugAnalysis}`,
+      config: {
+          systemInstruction: `You are a Solution Architect. 
+Based on the debugger's analysis, propose a minimal and effective structural fix.
+
+Instructions:
+- Specify which nodes to add, remove, or modify.
+- Specify which edges to reconnect.
+- Ensure the fix doesn't introduce new errors.
+
+Output the plan in clear, actionable steps.`
+      }
     });
     fixPlan = response.text.trim();
   } catch (e) {
     fixPlan = "Apply necessary structural corrections.";
   }
 
+  // --- Step 2.5: Critic (New Step) ---
+  onStatusUpdate('critic', 'Critic Agent is reviewing the fix plan...');
+  
+  let fixCritique = "";
+  try {
+      const response = await ai.models.generateContent({
+          model: model,
+          contents: `Error: "${errorMsg}"\nAnalysis: ${debugAnalysis}\nProposed Fix: ${fixPlan}`,
+          config: {
+              systemInstruction: `You are a Senior Reviewer. 
+Review the proposed fix plan. Does it actually solve the root cause? 
+Does it introduce new shape mismatches? 
+Is it the most efficient way to fix the error?
+You have access to Google Search to verify SOTA fixes for similar issues.
+
+Output a brief critique or "Approved" if perfect.`,
+              tools: [{ googleSearch: {} }]
+          }
+      });
+      fixCritique = response.text.trim();
+  } catch (e) {
+      fixCritique = "Approved";
+  }
+
   // --- Step 3: Patcher ---
   onStatusUpdate('patcher', 'Patcher Agent is applying the fix to the graph...');
   
-  const patcherPrompt = `
-    Role: DevOps Engineer.
-    Task: Apply the fix to the graph JSON.
-    
-    Current Graph: ${graphJson}
-    Fix Plan: ${fixPlan}
-    
-    Requirements:
-    1. Return the complete, valid JSON with "nodes" and "edges".
-    2. Maintain existing node positions where possible, offset new nodes if added.
-    3. Ensure all LayerTypes are valid.
-    4. Return ONLY raw JSON.
-  `;
-
   try {
     const response = await ai.models.generateContent({
-      model: MODEL_NAME,
-      contents: patcherPrompt,
-      config: { responseMimeType: "application/json" }
+      model: model,
+      contents: `Current Graph: ${graphJson}\nFix Plan: ${fixPlan}\nCritique: ${fixCritique}`,
+      config: { 
+          systemInstruction: `You are a Senior Patcher. 
+Apply the proposed fix plan (considering the critique) to the graph JSON.
+
+Requirements:
+1. Return a valid JSON object with "nodes" and "edges".
+2. Use ONLY valid LayerTypes.
+3. Maintain existing node IDs and positions where possible.
+4. Ensure all new nodes have unique IDs.
+
+Return ONLY the final JSON object.`,
+          // REMOVED responseMimeType and responseSchema as they conflict with previous tool use in history
+      }
     });
-    const finalJson = JSON.parse(response.text.trim());
+    const finalJson = JSON.parse(cleanJsonString(response.text));
     onStatusUpdate('complete', 'Fix applied successfully!');
     return sanitizeGraph(finalJson);
   } catch (e) {
     throw new Error("Patcher agent failed to generate valid JSON.");
   }
-};
+};