lila/data-pipeline/gemini.ts

import {
  DIFFICULTY_LEVELS,
  NOUN_GENDERS,
  type SupportedLanguageCode,
  type SupportedPos,
} from "@lila/shared";

export const DEFAULT_MODEL = "gemini-3.6-flash";

const API_BASE = "https://generativelanguage.googleapis.com/v1beta/models";
const MAX_ATTEMPTS = 5;
const RETRYABLE_STATUS = new Set([429, 500, 503]);

type ResponseSchema = Record<string, unknown>;

/**
 * OpenAPI-subset schema for Gemini structured output: an array of word
 * entries matching design-doc §6.3, with enums narrowed to this batch's
 * source/POS/target languages.
 */
export const buildEntriesResponseSchema = (
  sourceLanguage: SupportedLanguageCode,
  pos: SupportedPos,
  targetLanguages: readonly SupportedLanguageCode[],
): ResponseSchema => ({
  type: "ARRAY",
  items: {
    type: "OBJECT",
    properties: {
      headword: { type: "STRING" },
      language: { type: "STRING", enum: [sourceLanguage] },
      pos: { type: "STRING", enum: [pos] },
      senses: {
        type: "ARRAY",
        items: {
          type: "OBJECT",
          properties: {
            sense_index: { type: "INTEGER" },
            difficulty: { type: "STRING", enum: [...DIFFICULTY_LEVELS] },
            definitions: { type: "ARRAY", items: { type: "STRING" } },
            examples: { type: "ARRAY", items: { type: "STRING" } },
            translations: {
              type: "ARRAY",
              items: {
                type: "OBJECT",
                properties: {
                  target_language: {
                    type: "STRING",
                    enum: [...targetLanguages],
                  },
                  word: { type: "STRING" },
                  gender: {
                    type: "STRING",
                    enum: [...NOUN_GENDERS],
                    nullable: true,
                  },
                  difficulty: { type: "STRING", enum: [...DIFFICULTY_LEVELS] },
                },
                required: ["target_language", "word", "gender", "difficulty"],
              },
            },
          },
          required: [
            "sense_index",
            "difficulty",
            "definitions",
            "examples",
            "translations",
          ],
        },
      },
    },
    required: ["headword", "language", "pos", "senses"],
  },
});

const sleep = (ms: number): Promise<void> =>
  new Promise((resolve) => setTimeout(resolve, ms));

const retryDelayMs = (body: string, attempt: number): number => {
  const match = body.match(/"retryDelay":\s*"(\d+(?:\.\d+)?)s"/);
  if (match?.[1] !== undefined) {
    return Math.ceil(Number(match[1]) * 1000) + 500;
  }
  return 2 ** attempt * 2000;
};

const extractText = (body: unknown): string => {
  if (isRecord(body)) {
    const candidates = body["candidates"];
    if (Array.isArray(candidates) && isRecord(candidates[0])) {
      const candidate = candidates[0];
      const finishReason = candidate["finishReason"];
      if (finishReason !== undefined && finishReason !== "STOP") {
        throw new Error(
          `Gemini stopped early: finishReason=${JSON.stringify(finishReason)}`,
        );
      }
      const content = candidate["content"];
      if (isRecord(content)) {
        const parts = content["parts"];
        if (Array.isArray(parts) && isRecord(parts[0])) {
          const text = parts[0]["text"];
          if (typeof text === "string") return text;
        }
      }
    }
  }
  throw new Error("Gemini response contained no text candidate");
};

const isRecord = (value: unknown): value is Record<string, unknown> =>
  typeof value === "object" && value !== null && !Array.isArray(value);

export const generateContent = async (
  apiKey: string,
  model: string,
  prompt: string,
  responseSchema: ResponseSchema,
): Promise<string> => {
  let lastError = "";
  for (let attempt = 0; attempt < MAX_ATTEMPTS; attempt++) {
    let response: Response;
    try {
      response = await fetch(`${API_BASE}/${model}:generateContent`, {
        method: "POST",
        headers: {
          "Content-Type": "application/json",
          "x-goog-api-key": apiKey,
        },
        body: JSON.stringify({
          contents: [{ role: "user", parts: [{ text: prompt }] }],
          generationConfig: {
            responseMimeType: "application/json",
            responseSchema,
            temperature: 0.2,
          },
        }),
      });
    } catch (error) {
      lastError = `network error: ${error instanceof Error ? error.message : String(error)}`;
      await sleep(2 ** attempt * 2000);
      continue;
    }

    if (response.ok) {
      return extractText(await response.json());
    }

    const body = await response.text();
    lastError = `HTTP ${response.status}: ${body.slice(0, 500)}`;
    if (!RETRYABLE_STATUS.has(response.status)) {
      break;
    }
    await sleep(retryDelayMs(body, attempt));
  }
  throw new Error(`Gemini request failed after retries — ${lastError}`);
};

/** Parse the model's JSON text into an array of unknown entries. */
export const parseEntries = (rawText: string): unknown[] => {
  const stripped = rawText
    .trim()
    .replace(/^```(?:json)?\s*/i, "")
    .replace(/\s*```$/, "");
  const parsed: unknown = JSON.parse(stripped);
  if (!Array.isArray(parsed)) {
    throw new Error("Gemini response is not a JSON array");
  }
  return parsed;
};