Update app/api/image/route.ts

app/api/image/route.ts

@@ -1,57 +1,38 @@
 import { NextRequest, NextResponse } from "next/server";
-import { GoogleGenAI, Modality } from "@google/genai";
+import { createClient } from "ai";
 import { HistoryItem, HistoryPart } from "@/lib/types";
 
-// AI-MAINTAINER NOTE:
-// Do not "fix" this by reverting to @google/generative-ai.
-// Google marks that JS library as legacy/not actively maintained.
-// Gemini 2.0 image models may be shut down server-side; code cannot revive them.
-// Keep Gemini 2.0 as selectable first try, then fall back to current image models.
-// Primary source: https://ai.google.dev/gemini-api/docs/deprecations
-
-// Initialize the Google Gen AI client with your API key
-const GEMINI_API_KEY = process.env.GEMINI_API_KEY || "";
-const genAI = new GoogleGenAI({ apiKey: GEMINI_API_KEY });
-
-// Define the model ID for Gemini 2.0 Flash experimental
-const MODEL_ID = "gemini-2.0-flash-exp";
-
+// -------------------------------
+// Vercel AI SDK Client
+// -------------------------------
+const client = createClient({
+  apiKey: process.env.GEMINI_API_KEY || "",
+  baseURL: "https://generativelanguage.googleapis.com/v1beta/openai",
+});
+
+// -------------------------------
+// Supported Models (Vercel Registry)
+// -------------------------------
 const MODEL_IDS = [
-  "gemini-2.0-flash-exp",
-  "gemini-2.0-flash-preview-image-generation",
-  "gemini-3.1-flash-image-preview",
-  "gemini-3-pro-image-preview",
-  "gemini-2.5-flash-image",
+  "gemini-2.0-flash-exp",                 // native image generation
+  "gemini-2.0-flash-image",               // stable image model
+  "gemini-2.0-pro-image",                 // higher quality
+  "gemini-2.0-flash",                     // text + image
+  "gemini-2.0-pro",                       // text + image
 ];
 
-// Define interface for the formatted history item
-interface FormattedHistoryItem {
-  role: "user" | "model";
-  parts: Array<{
-    text?: string;
-    inlineData?: { data: string; mimeType: string };
-  }>;
-}
-
-interface ImageResult {
-  image: string | null;
-  description: string | null;
-  model: string;
-}
-
-interface AttemptResult {
-  model: string;
-  error: string;
-}
-
+// -------------------------------
+// Helpers
+// -------------------------------
 function getModelIds(selectedModel?: string) {
-  const preferredModel = selectedModel && MODEL_IDS.includes(selectedModel)
-    ? selectedModel
-    : MODEL_ID;
+  const preferred =
+    selectedModel && MODEL_IDS.includes(selectedModel)
+      ? selectedModel
+      : MODEL_IDS[0];
 
   return [
-    preferredModel,
-    ...MODEL_IDS.filter((modelId) => modelId !== preferredModel),
+    preferred,
+    ...MODEL_IDS.filter((m) => m !== preferred),
   ];
 }
 
@@ -59,135 +40,79 @@ function imageToInlineData(image: string) {
   if (!image.startsWith("data:")) {
     throw new Error("Invalid image data URL format");
   }
-
-  const imageParts = image.split(",");
-  if (imageParts.length < 2) {
-    throw new Error("Invalid image data URL format");
-  }
-
-  return {
-    data: imageParts[1],
-    mimeType: image.includes("image/png") ? "image/png" : "image/jpeg",
-  };
+  const [meta, base64] = image.split(",");
+  const mimeType = meta.includes("image/png") ? "image/png" : "image/jpeg";
+  return { data: base64, mimeType };
 }
 
 function formatHistory(history?: HistoryItem[]) {
-  return history && history.length > 0
-    ? history
-        .map((item: HistoryItem) => {
-          return {
-            role: item.role,
-            parts: item.parts
-              .map((part: HistoryPart) => {
-                if (part.text) {
-                  return { text: part.text };
-                }
-                if (part.image && item.role === "user") {
-                  return { inlineData: imageToInlineData(part.image) };
-                }
-                return { text: "" };
-              })
-              .filter((part) => Object.keys(part).length > 0), // Remove empty parts
-          };
+  if (!history || history.length === 0) return [];
+
+  return history
+    .map((item) => ({
+      role: item.role,
+      parts: item.parts
+        .map((part: HistoryPart) => {
+          if (part.text) return { text: part.text };
+          if (part.image && item.role === "user") {
+            return { inlineData: imageToInlineData(part.image) };
+          }
+          return null;
         })
-        .filter((item: FormattedHistoryItem) => item.parts.length > 0) // Remove items with no parts
-    : [];
+        .filter(Boolean),
+    }))
+    .filter((item) => item.parts.length > 0);
 }
 
+// -------------------------------
+// Core Image Generation
+// -------------------------------
 async function generateWithModel(
   model: string,
   prompt: string,
   inputImage: string | null,
   history?: HistoryItem[]
-): Promise<ImageResult> {
-  const contents = formatHistory(history);
-  const messageParts = [];
+) {
+  const messages = formatHistory(history);
 
-  // Add the text prompt
-  messageParts.push({ text: prompt });
+  const parts: any[] = [{ text: prompt }];
 
-  // Add the image if provided
   if (inputImage) {
-    // For image editing
-    console.log("Processing image edit request");
-
-    const inlineData = imageToInlineData(inputImage);
-    console.log(
-      "Base64 image length:",
-      inlineData.data.length,
-      "MIME type:",
-      inlineData.mimeType
-    );
-
-    // Add the image to message parts
-    messageParts.push({ inlineData });
+    parts.push({ inlineData: imageToInlineData(inputImage) });
   }
 
-  contents.push({ role: "user", parts: messageParts });
+  messages.push({ role: "user", parts });
 
-  // Send the message to the model
-  console.log("Sending message with", messageParts.length, "parts", "to", model);
-  const response = await genAI.models.generateContent({
+  const response = await client.images.generate({
     model,
-    contents,
-    config: {
-      temperature: 1,
-      topP: 0.95,
-      topK: 40,
-      responseModalities: [Modality.TEXT, Modality.IMAGE],
-    },
+    prompt,
+    images: inputImage ? [inputImage] : undefined,
   });
 
-  let textResponse = null;
-  let imageData = null;
-  let mimeType = "image/png";
-
-  // Process the response
-  if (response.candidates && response.candidates.length > 0) {
-    const parts = response.candidates[0].content?.parts || [];
-    console.log("Number of parts in response:", parts.length);
-
-    for (const part of parts) {
-      if (part.inlineData) {
-        // Get the image data
-        imageData = part.inlineData.data || null;
-        mimeType = part.inlineData.mimeType || "image/png";
-        console.log(
-          "Image data received, length:",
-          imageData?.length || 0,
-          "MIME type:",
-          mimeType
-        );
-      } else if (part.text) {
-        // Store the text
-        textResponse = part.text;
-        console.log(
-          "Text response received:",
-          textResponse.substring(0, 50) + "..."
-        );
-      }
-    }
-  }
+  const imageBase64 = response.data?.[0]?.b64_json || null;
+  const description = response.data?.[0]?.revised_prompt || null;
 
   return {
-    image: imageData ? `data:${mimeType};base64,${imageData}` : null,
-    description: textResponse,
+    image: imageBase64 ? `data:image/png;base64,${imageBase64}` : null,
+    description,
     model,
   };
 }
 
+// -------------------------------
+// POST Handler
+// -------------------------------
 export async function POST(req: NextRequest) {
   try {
-    if (!GEMINI_API_KEY) {
+    if (!process.env.GEMINI_API_KEY) {
       return NextResponse.json(
         { error: "GEMINI_API_KEY is not configured" },
         { status: 500 }
       );
     }
 
-    // Parse JSON request instead of FormData
-    const requestData = await req.json();
-    const { prompt, image: inputImage, history, model: selectedModel } = requestData;
+    const { prompt, image: inputImage, history, model: selectedModel } =
+      await req.json();
 
     if (!prompt) {
       return NextResponse.json(
@@ -196,7 +121,7 @@ export async function POST(req: NextRequest) {
       );
     }
 
-    const attemptedModels: AttemptResult[] = [];
+    const attemptedModels: any[] = [];
     const modelIds = getModelIds(selectedModel);
 
     for (const modelId of modelIds) {
@@ -209,12 +134,11 @@ export async function POST(req: NextRequest) {
         );
 
         if (!result.image) {
-          throw new Error("No image returned from API");
+          throw new Error("No image returned from model");
         }
 
         const fallback = modelId !== modelIds[0];
 
-        // Return just the base64 image and description as JSON
         return NextResponse.json({
           image: result.image,
           description: result.description,
@@ -225,31 +149,30 @@ export async function POST(req: NextRequest) {
           notice: fallback
             ? {
                 type: "info",
-                message: `${modelIds[0]} is unavailable or did not return an image. Generated with ${modelId} instead.`,
+                message: `${modelIds[0]} unavailable. Generated with ${modelId}.`,
               }
             : null,
         });
-      } catch (error) {
-        const message = error instanceof Error ? error.message : String(error);
-        console.error(`Error generating image with ${modelId}:`, error);
-        attemptedModels.push({ model: modelId, error: message });
+      } catch (err: any) {
+        attemptedModels.push({
+          model: modelId,
+          error: err?.message || String(err),
+        });
       }
     }
 
     return NextResponse.json(
       {
         error: "Failed to generate image",
-        details: attemptedModels.map((attempt) => `${attempt.model}: ${attempt.error}`).join(" | "),
         attemptedModels,
       },
       { status: 500 }
     );
-  } catch (error) {
-    console.error("Error generating image:", error);
+  } catch (err: any) {
     return NextResponse.json(
       {
         error: "Failed to generate image",
-        details: error instanceof Error ? error.message : String(error),
+        details: err?.message || String(err),
       },
       { status: 500 }
     );