Update services/geminiService.ts

services/geminiService.ts

@@ -1,3 +1,4 @@
+
 import { GoogleGenAI } from "@google/genai";
 import { Node, Edge } from 'reactflow';
 import { NodeData, LayerType } from '../types';
@@ -44,8 +45,11 @@ const buildRawPrompt = (nodes: Node<NodeData>[], edges: Edge[]): string => {
           if (typeof v === 'object' && v !== null) {
               return `${k}=${JSON.stringify(v)}`;
           }
+          // Don't clutter standard output with huge code blocks, handle separately
+          if (k === 'definition_code' || k === 'imports') return null;
           return `${k}=${v}`;
       })
+      .filter(p => p !== null)
       .join(', ');
       
     rawSpec += `- [ID: ${node.id}] TYPE: ${node.data.type} | LABEL: ${node.data.label}\n`;
@@ -56,6 +60,12 @@ const buildRawPrompt = (nodes: Node<NodeData>[], edges: Edge[]): string => {
     // Specific instruction for Custom Layers
     if (node.data.type === LayerType.CUSTOM) {
         rawSpec += `  CUSTOM_NOTE: Instantiate using class '${node.data.params.class_name}' with args '${node.data.params.args}'.\n`;
+        if (node.data.params.imports) {
+             rawSpec += `  CUSTOM_IMPORTS: ${node.data.params.imports}\n`;
+        }
+        if (node.data.params.definition_code) {
+             rawSpec += `  CUSTOM_CODE_DEFINITION:\n${node.data.params.definition_code}\n`;
+        }
     }
   });
   
@@ -100,8 +110,9 @@ export const generateRefinedPrompt = async (nodes: Node<NodeData>[], edges: Edge
        - 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. Do NOT write the Python code yourself. Write the PROMPT that asks for the code.
-    7. Ensure the tone is technical and precise.
+    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.
     `;