Update services/geminiService.ts

services/geminiService.ts

@@ -39,6 +39,8 @@ const evaluationSchema = {
 const graphResponseSchema = {
   type: Type.OBJECT,
   properties: {
+    model_type: { type: Type.STRING, description: "The overarching type of the model (e.g., 'cnn_transformer_hybrid', 'vision_transformer')." },
+    input_shape: { type: Type.ARRAY, items: { type: Type.NUMBER }, description: "Input shape [C,H,W] or [N,C,H,W] or other." },
     reasoning: { 
         type: Type.STRING, 
         description: "Brief explanation of architectural choices or fixes applied." 
@@ -591,6 +593,7 @@ export const applyPatches = (code: string, patches: { search: string, replace: s
  */
 export const evaluateCode = async (
     prompt: string,
+    irRepresentation: string,
     code: string,
     model: string = DEFAULT_MODEL
 ): Promise<{ score: number, missing_components: string[], bugs: string[], improvement_tasks: string[] }> => {
@@ -598,10 +601,10 @@ export const evaluateCode = async (
     try {
         const response = await ai.models.generateContent({
             model: model,
-            contents: `Original Prompt: "${prompt}"\n\nGenerated Code:\n${code}`,
+            contents: `Original Prompt: "${prompt}"\n\nArchitecture IR:\n${irRepresentation}\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.
+                systemInstruction: `You are a Senior QA Engineer. Compare the generated PyTorch code against the original prompt and the formal IR schema.
+Identify any missing requirements, shape mismatches, logical bugs, or incomplete components.
 Generate a list of targeted improvement tasks.
 Return the evaluation in JSON format.`,
                 responseMimeType: "application/json",
@@ -656,14 +659,18 @@ Return the patches in JSON format.`,
  * Uses the detailed prompt to write, review, and polish PyTorch code.
  */
 export const generateCodeWithAgents = async (
+    nodes: Node[],
+    edges: Edge[],
     promptText: string,
     onStatusUpdate: (status: AgentStatus, log: string) => void,
     model: string = DEFAULT_MODEL
 ): Promise<{ code: string, evaluation?: any, initialPatches?: any }> => {
   const ai = getAiClient();
+  
+  const irRepresentation = JSON.stringify({ nodes, edges });
 
   // --- Step 1: Coder ---
-  onStatusUpdate('architect', 'Coder Agent is writing initial PyTorch implementation...');
+  onStatusUpdate('architect', 'Coder Agent is writing initial PyTorch implementation from IR...');
 
   const coderSystemInstruction = `You are a Senior Deep Learning Engineer. 
 Your task is to write clean, modular, and production-ready PyTorch code.
@@ -673,11 +680,11 @@ Coding Standards:
 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.
+5. Include a robust 'if __name__ == "__main__":' block that acts as a HARD VERIFIER. It must instantiate the model, run a dummy forward pass with torch.randn, print output shapes, calculate parameter counts, and print a 'Verification Passed!' message.
 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.
+   - CRITICAL: If using SiLU (nn.SiLU) as the activation, use nonlinearity='linear' (gain=1).
 
 Return ONLY the Python code. No markdown backticks.`;
 
@@ -685,13 +692,13 @@ Return ONLY the Python code. No markdown backticks.`;
   try {
       const response = await ai.models.generateContent({
         model: model,
-        contents: `Request: "${promptText}"`,
+        contents: `Original Request: "${promptText}"\n\nArchitecture IR:\n${irRepresentation}`,
         config: { 
             systemInstruction: coderSystemInstruction,
             tools: [{ googleSearch: {} }, { urlContext: {} }]
         }
       });
-      draftCode = response.text.trim().replace(/```python/g, '').replace(/```/g, '');
+      draftCode = cleanJsonString(response.text);
   } catch(e) {
       throw e;
   }
@@ -700,12 +707,13 @@ Return ONLY the Python code. No markdown backticks.`;
   onStatusUpdate('critic', 'Reviewer Agent is analyzing code for bugs and optimization...');
 
   const reviewerSystemInstruction = `You are a Senior Code Reviewer. 
-Analyze the provided PyTorch code for:
+Analyze the provided PyTorch code against the Architecture IR constraints:
+- Shape mismatch according to the IR
 - 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.
+- Adherence to the original architecture request and IR graph.
 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".`;
@@ -715,7 +723,7 @@ Provide a concise, technical critique. If the code is perfect, say "No changes n
       const response = await ai.models.generateContent({
         model: model,
         contents: [
-            { role: 'user', parts: [{ text: `Original Request: "${promptText}"` }] },
+            { role: 'user', parts: [{ text: `Architecture IR:\n${irRepresentation}\n\nOriginal Request: "${promptText}"` }] },
             { role: 'model', parts: [{ text: draftCode }] },
             { role: 'user', parts: [{ text: "Please review this code." }] }
         ],
@@ -747,7 +755,7 @@ Return ONLY the final Python code. No markdown backticks.`;
       const response = await ai.models.generateContent({
         model: model,
         contents: [
-            { role: 'user', parts: [{ text: `Original Request: "${promptText}"` }] },
+            { role: 'user', parts: [{ text: `Architecture IR:\n${irRepresentation}\n\nOriginal Request: "${promptText}"` }] },
             { role: 'model', parts: [{ text: draftCode }] },
             { role: 'user', parts: [{ text: `Critique: ${critique}\n\nApply the critique and provide the final polished code.` }] }
         ],
@@ -757,7 +765,7 @@ Return ONLY the final Python code. No markdown backticks.`;
       
       // --- Step 4: Evaluator (New Step) ---
       onStatusUpdate('evaluator', 'Evaluator Agent is verifying requirements...');
-      const evaluation = await evaluateCode(promptText, finalCode, model);
+      const evaluation = await evaluateCode(promptText, irRepresentation, finalCode, model);
       
       let initialPatches = null;
       if (evaluation.score < 90 && evaluation.improvement_tasks.length > 0) {