Upload geminiService.ts

services/geminiService.ts

@@ -0,0 +1,468 @@
+import { GoogleGenAI } from "@google/genai";
+import { Node, Edge } from 'reactflow';
+import { NodeData, LayerType } from '../types';
+import { LAYER_DEFINITIONS } from '../constants';
+
+// Key Management
+let userApiKey = typeof window !== 'undefined' ? localStorage.getItem('gemini_api_key') || '' : '';
+const defaultEnvKey = process.env.API_KEY || '';
+
+export const setUserApiKey = (key: string) => {
+  userApiKey = key;
+  if (typeof window !== 'undefined') {
+    localStorage.setItem('gemini_api_key', key);
+  }
+};
+
+export const getUserApiKey = () => userApiKey || defaultEnvKey;
+
+const getAiClient = () => {
+  const key = getUserApiKey();
+  if (!key) throw new Error("API Key is missing. Please add your Gemini API Key.");
+  return new GoogleGenAI({ apiKey: key });
+};
+
+const MODEL_NAME = 'gemini-2.5-flash';
+
+export type AgentStatus = 'idle' | 'architect' | 'critic' | 'refiner' | 'complete' | 'error';
+
+/**
+ * Internal helper to build the raw specification string.
+ * This ensures all data is captured before sending to the AI for refinement.
+ */
+const buildRawPrompt = (nodes: Node<NodeData>[], edges: Edge[]): string => {
+  // Sort nodes by vertical position to list them in a logical flow
+  const sortedNodes = [...nodes].sort((a, b) => a.position.y - b.position.y);
+  
+  let rawSpec = "### Raw Architecture Data\n\n";
+  
+  rawSpec += "**1. Nodes (Layers):**\n";
+  sortedNodes.forEach(node => {
+    // Format parameters nicely, handling objects
+    const params = Object.entries(node.data.params)
+      .map(([k, v]) => {
+          if (typeof v === 'object' && v !== null) {
+              return `${k}=${JSON.stringify(v)}`;
+          }
+          return `${k}=${v}`;
+      })
+      .join(', ');
+      
+    rawSpec += `- [ID: ${node.id}] TYPE: ${node.data.type} | LABEL: ${node.data.label}\n`;
+    if (params) {
+      rawSpec += `  PARAMS: ${params}\n`;
+    }
+    
+    // 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`;
+    }
+  });
+  
+  rawSpec += "\n**2. Connectivity (Edges):**\n";
+  if (edges.length === 0) {
+    rawSpec += "- No connections defined.\n";
+  } else {
+    edges.forEach(edge => {
+       const sourceNode = nodes.find(n => n.id === edge.source);
+       const targetNode = nodes.find(n => n.id === edge.target);
+       const sourceName = sourceNode ? `${sourceNode.data.label} (ID:${sourceNode.id})` : edge.source;
+       const targetName = targetNode ? `${targetNode.data.label} (ID:${targetNode.id})` : edge.target;
+       
+       rawSpec += `- ${sourceName} -> ${targetName}\n`;
+    });
+  }
+  
+  return rawSpec;
+};
+
+/**
+ * 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> => {
+    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. Do NOT write the Python code yourself. Write the PROMPT that asks for the code.
+    7. 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,
+        });
+        return response.text.trim();
+    } catch (error) {
+        console.error("Prompt refinement failed:", error);
+        return buildRawPrompt(nodes, edges) + "\n\n(AI Refinement failed, showing raw spec)";
+    }
+};
+
+export const validateArchitecture = async (nodes: Node<NodeData>[], edges: Edge[]): Promise<string> => {
+   const ai = getAiClient();
+   const graphRepresentation = {
+    nodes: nodes.map(n => ({
+      id: n.id,
+      type: n.data.type,
+      parameters: n.data.params
+    })),
+    edges: edges.map(e => ({
+      source: e.source,
+      target: e.target
+    }))
+  };
+
+  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,
+    });
+    return response.text;
+  } catch (error) {
+    return "Error validating architecture.";
+  }
+}
+
+/**
+ * Gets AI suggestions for improving the architecture.
+ */
+export const getArchitectureSuggestions = async (nodes: Node<NodeData>[], edges: Edge[]): Promise<string> => {
+   const ai = getAiClient();
+   const graphRepresentation = {
+    nodes: nodes.map(n => ({
+      id: n.id,
+      type: n.data.type,
+      parameters: n.data.params,
+      label: n.data.label
+    })),
+    edges: edges.map(e => ({
+      source: e.source,
+      target: e.target
+    }))
+  };
+
+  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,
+    });
+    return response.text;
+  } catch (error) {
+    return "Error generating suggestions.";
+  }
+}
+
+/**
+ * Implements the suggestions automatically.
+ */
+export const implementArchitectureSuggestions = async (
+  nodes: Node<NodeData>[], 
+  edges: Edge[], 
+  suggestions: string
+): Promise<{ nodes: any[], edges: any[] }> => {
+  const ai = getAiClient();
+  
+  const graphRepresentation = {
+    nodes: nodes.map(n => ({
+      id: n.id,
+      type: n.data.type,
+      parameters: n.data.params,
+      label: n.data.label,
+      position: n.position
+    })),
+    edges: edges.map(e => ({
+      source: e.source,
+      target: e.target
+    }))
+  };
+
+  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" }
+    });
+    return JSON.parse(response.text.trim());
+  } catch (error) {
+    throw new Error("Failed to implement suggestions.");
+  }
+};
+
+
+/**
+ * Multi-agent graph generation.
+ * 1. Architect Agent: Drafts the layout.
+ * 2. Critic Agent: Reviews for errors.
+ * 3. Refiner Agent: Produces final JSON.
+ */
+export const generateGraphWithAgents = async (
+  userPrompt: string, 
+  currentNodes: Node<NodeData>[] = [],
+  onStatusUpdate: (status: AgentStatus, log: string) => void
+): Promise<{ nodes: any[], edges: any[] } | null> => {
+  const ai = getAiClient();
+
+  const layerSchema = Object.values(LAYER_DEFINITIONS).map(l => ({
+    type: l.type,
+    params: l.parameters.map(p => ({ name: p.name, type: p.type, options: p.options }))
+  }));
+  const schemaStr = JSON.stringify(layerSchema);
+
+  // --- Step 1: Architect ---
+  onStatusUpdate('architect', 'Architect is drafting initial layout...');
+  
+  const context = currentNodes.length > 0 
+    ? `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.
+  `;
+
+  let draftJsonStr = "";
+  try {
+    const response = await ai.models.generateContent({
+        model: MODEL_NAME,
+        contents: architectPrompt,
+        config: { responseMimeType: "application/json" }
+    });
+    draftJsonStr = response.text.trim();
+  } catch (e) {
+      throw new Error("Architect agent failed to generate draft.");
+  }
+
+  // --- 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".
+  `;
+
+  let critique = "";
+  try {
+     const response = await ai.models.generateContent({
+        model: MODEL_NAME,
+        contents: criticPrompt,
+    });
+    critique = response.text.trim();
+  } catch (e) {
+      console.warn("Critic agent failed, proceeding with draft.");
+      critique = "No critique available.";
+  }
+
+  // --- 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.
+  `;
+
+  try {
+    const response = await ai.models.generateContent({
+        model: MODEL_NAME,
+        contents: refinerPrompt,
+        config: { responseMimeType: "application/json" }
+    });
+    const finalJson = JSON.parse(response.text.trim());
+    onStatusUpdate('complete', 'Architecture built successfully!');
+    return finalJson;
+  } catch (e) {
+     throw new Error("Refiner agent failed to parse final JSON.");
+  }
+};
+
+/**
+ * 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
+): Promise<string> => {
+  const ai = getAiClient();
+
+  // --- Step 1: Coder (Architect) ---
+  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.
+  `;
+
+  let draftCode = "";
+  try {
+      const response = await ai.models.generateContent({
+        model: MODEL_NAME,
+        contents: coderPrompt
+      });
+      draftCode = response.text.trim().replace(/```python/g, '').replace(/```/g, '');
+  } catch(e) {
+      throw new Error("Coder agent failed.");
+  }
+
+  // --- Step 2: Reviewer (Critic) ---
+  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".
+  `;
+  
+  let critique = "";
+  try {
+      const response = await ai.models.generateContent({
+        model: MODEL_NAME,
+        contents: reviewPrompt
+      });
+      critique = response.text.trim();
+  } catch(e) {
+      critique = "No critique available.";
+  }
+
+  // --- Step 3: Polisher (Refiner) ---
+  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.
+  `;
+
+  try {
+      const response = await ai.models.generateContent({
+        model: MODEL_NAME,
+        contents: polisherPrompt
+      });
+      let finalCode = response.text.trim().replace(/```python/g, '').replace(/```/g, '');
+      onStatusUpdate('complete', 'Code generation complete!');
+      return finalCode;
+  } catch(e) {
+      throw new Error("Polisher agent failed.");
+  }
+};