Update services/geminiService.ts

services/geminiService.ts

@@ -4,6 +4,38 @@ import { NodeData, LayerType } from '../types';
 import { LAYER_DEFINITIONS } from '../constants';
 
 // --- SCHEMAS ---
+const patchSchema = {
+  type: Type.OBJECT,
+  properties: {
+    reasoning: { type: Type.STRING, description: "Why these changes are being made." },
+    summary: { type: Type.STRING, description: "A concise summary of what was updated (e.g., 'Added dropout to encoder')." },
+    files_changed: { type: Type.ARRAY, items: { type: Type.STRING }, description: "List of files modified." },
+    patches: {
+      type: Type.ARRAY,
+      items: {
+        type: Type.OBJECT,
+        properties: {
+          search: { type: Type.STRING, description: "The exact code block to find. Must be unique." },
+          replace: { type: Type.STRING, description: "The code to replace it with." }
+        },
+        required: ["search", "replace"]
+      }
+    }
+  },
+  required: ["patches", "summary", "files_changed"]
+};
+
+const evaluationSchema = {
+  type: Type.OBJECT,
+  properties: {
+    score: { type: Type.NUMBER, description: "0-100 score of how well requirements were met." },
+    missing_components: { type: Type.ARRAY, items: { type: Type.STRING } },
+    bugs: { type: Type.ARRAY, items: { type: Type.STRING } },
+    improvement_tasks: { type: Type.ARRAY, items: { type: Type.STRING } }
+  },
+  required: ["score", "improvement_tasks"]
+};
+
 const graphResponseSchema = {
   type: Type.OBJECT,
   properties: {
@@ -94,7 +126,7 @@ const getAiClient = () => {
 
 export const DEFAULT_MODEL = 'gemini-2.5-flash';
 
-export type AgentStatus = 'idle' | 'architect' | 'critic' | 'refiner' | 'debugger' | 'patcher' | 'complete' | 'error';
+export type AgentStatus = 'idle' | 'architect' | 'critic' | 'refiner' | 'evaluator' | 'debugger' | 'patcher' | 'complete' | 'error';
 
 /**
  * Internal helper to build the raw specification string.
@@ -539,15 +571,95 @@ Return ONLY the final JSON object.`;
   }
 };
 
+/**
+ * Helper: Applies a list of patches to a code string.
+ */
+export const applyPatches = (code: string, patches: { search: string, replace: string }[]): string => {
+    let updatedCode = code;
+    for (const patch of patches) {
+        if (updatedCode.includes(patch.search)) {
+            updatedCode = updatedCode.replace(patch.search, patch.replace);
+        } else {
+            console.warn("Patch failed: Search block not found.", patch.search);
+        }
+    }
+    return updatedCode;
+};
+
+/**
+ * Evaluates generated code against the original prompt.
+ */
+export const evaluateCode = async (
+    prompt: string,
+    code: string,
+    model: string = DEFAULT_MODEL
+): Promise<{ score: number, missing_components: string[], bugs: string[], improvement_tasks: string[] }> => {
+    const ai = getAiClient();
+    try {
+        const response = await ai.models.generateContent({
+            model: model,
+            contents: `Original Prompt: "${prompt}"\n\nGenerated Code:\n${code}`,
+            config: {
+                systemInstruction: `You are a Senior QA Engineer. Compare the generated PyTorch code against the original prompt.
+Identify any missing requirements, logical bugs, or incomplete components.
+Generate a list of targeted improvement tasks.
+Return the evaluation in JSON format.`,
+                responseMimeType: "application/json",
+                responseSchema: evaluationSchema
+            }
+        });
+        return JSON.parse(cleanJsonString(response.text));
+    } catch (error) {
+        console.error("Evaluation failed:", error);
+        return { score: 0, missing_components: [], bugs: [], improvement_tasks: ["Manual review required."] };
+    }
+};
+
+/**
+ * Generates patches for code based on instructions.
+ */
+export const generateCodePatches = async (
+    code: string,
+    instructions: string,
+    model: string = DEFAULT_MODEL
+): Promise<{ reasoning: string, summary: string, files_changed: string[], patches: { search: string, replace: string }[] }> => {
+    const ai = getAiClient();
+    try {
+        const response = await ai.models.generateContent({
+            model: model,
+            contents: `Current Code:\n${code}\n\nInstructions for Modification:\n${instructions}`,
+            config: {
+                systemInstruction: `You are a Principal Software Engineer. Instead of rewriting the entire file, you provide targeted patches.
+A patch consists of a 'search' block (the exact code to find) and a 'replace' block (the new code).
+The 'search' block MUST be a unique, contiguous substring from the current code.
+Include all necessary imports if they are missing.
+
+Output Requirements:
+1. 'summary': A concise one-line description of the update.
+2. 'files_changed': An array of filenames (e.g., ["model.py"]).
+3. 'patches': The search/replace pairs.
+
+Return the patches in JSON format.`,
+                responseMimeType: "application/json",
+                responseSchema: patchSchema
+            }
+        });
+        return JSON.parse(cleanJsonString(response.text));
+    } catch (error) {
+        console.error("Patch generation failed:", error);
+        throw error;
+    }
+};
+
 /**
  * Multi-agent code generation.
  * Uses the detailed prompt to write, review, and polish PyTorch code.
  */
 export const generateCodeWithAgents = async (
-  promptText: string,
-  onStatusUpdate: (status: AgentStatus, log: string) => void,
-  model: string = DEFAULT_MODEL
-): Promise<string> => {
+    promptText: string,
+    onStatusUpdate: (status: AgentStatus, log: string) => void,
+    model: string = DEFAULT_MODEL
+): Promise<{ code: string, evaluation?: any, initialPatches?: any }> => {
   const ai = getAiClient();
 
   // --- Step 1: Coder ---
@@ -563,6 +675,9 @@ Coding Standards:
 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.
+7. Initialization Strategy:
+   - Use Kaiming initialization (nn.init.kaiming_normal_) for Conv2d and Linear layers.
+   - CRITICAL: If using SiLU (nn.SiLU) as the activation, use nonlinearity='linear' (gain=1) in the Kaiming initialization instead of 'relu'. This better preserves feature variance for SiLU.
 
 Return ONLY the Python code. No markdown backticks.`;
 
@@ -639,8 +754,21 @@ Return ONLY the final Python code. No markdown backticks.`;
         config: { systemInstruction: polisherSystemInstruction }
       });
       let finalCode = cleanJsonString(response.text);
+      
+      // --- Step 4: Evaluator (New Step) ---
+      onStatusUpdate('evaluator', 'Evaluator Agent is verifying requirements...');
+      const evaluation = await evaluateCode(promptText, finalCode, model);
+      
+      let initialPatches = null;
+      if (evaluation.score < 90 && evaluation.improvement_tasks.length > 0) {
+          onStatusUpdate('refiner', 'Applying automated improvements based on evaluation...');
+          const patchResult = await generateCodePatches(finalCode, evaluation.improvement_tasks.join("\n"), model);
+          finalCode = applyPatches(finalCode, patchResult.patches);
+          initialPatches = patchResult;
+      }
+
       onStatusUpdate('complete', 'Code generation complete!');
-      return finalCode;
+      return { code: finalCode, evaluation, initialPatches };
   } catch(e) {
       throw new Error("Polisher agent failed to generate code.");
   }