karate_halloween/lib/ai/gemini.ts

import { ContentSafetyError } from "./safety";
export interface GeminiGenerateOptions {
  model: string;
  prompt: string;
  negativePrompt?: string;
  imageBase64: string;
  mimeType: string;
  timeoutMs: number;
  enableSfwCheck: boolean;
}

export interface GeminiResult {
  imageUrl: string;
  provider: string;
  model: string;
}

export async function generateWithGemini(
  options: GeminiGenerateOptions
): Promise<GeminiResult> {
  const apiKey = process.env.GEMINI_API_KEY;
  if (!apiKey) {
    throw new Error("GEMINI_API_KEY is not configured in environment variables");
  }

  // Clean base64 string
  const cleanBase64 = options.imageBase64.replace(/^data:image\/\w+;base64,/, "");

  const controller = new AbortController();
  const timeoutId = setTimeout(() => controller.abort(), options.timeoutMs);

  try {
    const url = `https://generativelanguage.googleapis.com/v1beta/models/${options.model}:generateContent?key=${apiKey}`;

    const requestBody: any = {
      contents: [
        {
          role: "user",
          parts: [
            {
              text: `${options.prompt}\n\nNegative directives: ${options.negativePrompt || "none"}`,
            },
            {
              inlineData: {
                mimeType: options.mimeType || "image/jpeg",
                data: cleanBase64,
              },
            },
          ],
        },
      ],
    };

    if (options.enableSfwCheck) {
      requestBody.safetySettings = [
        {
          category: "HARM_CATEGORY_SEXUALLY_EXPLICIT",
          threshold: "BLOCK_MEDIUM_AND_ABOVE",
        },
        {
          category: "HARM_CATEGORY_HATE_SPEECH",
          threshold: "BLOCK_MEDIUM_AND_ABOVE",
        },
        {
          category: "HARM_CATEGORY_HARASSMENT",
          threshold: "BLOCK_MEDIUM_AND_ABOVE",
        },
        {
          category: "HARM_CATEGORY_DANGEROUS_CONTENT",
          threshold: "BLOCK_MEDIUM_AND_ABOVE",
        },
      ];
    }

    const response = await fetch(url, {
      method: "POST",
      headers: {
        "Content-Type": "application/json",
      },
      body: JSON.stringify(requestBody),
      signal: controller.signal,
    });

    clearTimeout(timeoutId);

    if (!response.ok) {
      const errorText = await response.text();
      if (
        errorText.includes("PROHIBITED_CONTENT") ||
        errorText.includes("HARM_CATEGORY") ||
        errorText.includes("SAFETY")
      ) {
        throw new ContentSafetyError(
          `Image rejected by content safety filter: PROHIBITED_CONTENT`
        );
      }
      throw new Error(`Gemini API error [${response.status}]: ${errorText}`);
    }

    const data = await response.json();

    // Check safety feedback
    if (data.promptFeedback?.blockReason) {
      throw new ContentSafetyError(
        `Image rejected by content safety filter: ${data.promptFeedback.blockReason}`
      );
    }

    // Extract generated image from candidates
    const candidate = data.candidates?.[0];
    if (!candidate) {
      throw new Error("Gemini returned no candidates in response");
    }

    // Check candidate finishReason for safety blocks
    if (candidate.finishReason === "SAFETY" || candidate.finishReason === "BLOCKLIST") {
      throw new ContentSafetyError(
        "Image rejected by content safety filter: output triggered Safe-for-Work (SFW) protection."
      );
    }

    // Look for inline image in response parts
    const parts = candidate.content?.parts || [];
    for (const part of parts) {
      if (part.inlineData?.data) {
        const mime = part.inlineData.mimeType || "image/png";
        return {
          imageUrl: `data:${mime};base64,${part.inlineData.data}`,
          provider: "gemini",
          model: options.model,
        };
      }
    }

    // If text returned or URL in text
    const textPart = parts.find((p: any) => p.text)?.text || "";
    const urlMatch = textPart.match(/https?:\/\/[^\s"']+\.(?:png|jpg|jpeg|webp)/i);
    if (urlMatch) {
      return {
        imageUrl: urlMatch[0],
        provider: "gemini",
        model: options.model,
      };
    }

    throw new Error(
      "Gemini did not return an image part in the response payload. Please verify model image generation capabilities."
    );
  } catch (error: any) {
    clearTimeout(timeoutId);
    if (error.name === "AbortError") {
      throw new Error(`Gemini request timed out after ${options.timeoutMs}ms`);
    }
    throw error;
  }
}