lila/data-pipeline/utils/enrich-word.ts

import { buildSystemPrompt } from "../config/prompt.js";
import { createAdapter } from "./llm-adapters/factory.js";
import { BATCH_CONFIG } from "../config/batch.js";
import type { Language, Pos, EnrichedSense } from "./merge-enriched-data.js";
import { LANG_MAP, POS_MAP, ALL_LANGUAGES } from "../config/constants.js";

interface LlmResponse {
  content: string;
  promptTokens: number;
  completionTokens: number;
  totalTokens: number;
  promptTimeMs: number;
  completionTimeMs: number;
}

export interface EnrichmentResult {
  results: Map<string, EnrichedSense[]>;
  metrics: {
    promptTokens: number;
    completionTokens: number;
    totalTokens: number;
    promptTimeMs: number;
    completionTimeMs: number;
  };
}

/**
 * Calls the LLM with the enrichment prompt.
 * Returns the response content and timing metrics.
 */
async function callLlm(
  words: string[],
  rawLanguage: string,
  rawPos: string,
): Promise<LlmResponse> {
  const adapter = createAdapter();
  const sourceCode = LANG_MAP[rawLanguage] || rawLanguage;
  const targetLanguages = ALL_LANGUAGES.filter((lang) => lang !== sourceCode);
  const prompt = buildSystemPrompt(rawLanguage, rawPos, targetLanguages);
  return adapter.call(words, prompt);
}

/**
 * Strips markdown code blocks and extracts the JSON object from raw LLM output.
 * Throws if no valid JSON object braces are found.
 */
function sanitizeLlmOutput(raw: string): string {
  const cleaned = raw.replace(/```json\s*/g, "").replace(/```\s*$/g, "");
  const start = cleaned.indexOf("{");
  const end = cleaned.lastIndexOf("}");
  if (start === -1 || end === -1) {
    throw new Error("No JSON object found in LLM output");
  }
  return cleaned.slice(start, end + 1);
}

/**
 * Parses the LLM response string into a JavaScript object.
 * Throws if the response is not valid JSON or not an object with expected keys.
 */
export function parseLlmResponse(
  rawJson: string,
  expectedWords: string[],
): Record<string, unknown> {
  let parsed: unknown;

  try {
    const sanitized = sanitizeLlmOutput(rawJson);
    parsed = JSON.parse(sanitized);
  } catch (error: unknown) {
    throw new Error(`Failed to parse LLM output as JSON: ${rawJson}`, {
      cause: error,
    });
  }

  if (typeof parsed !== "object" || parsed === null || Array.isArray(parsed)) {
    throw new Error("LLM output is not a JSON object");
  }

  const obj = parsed as Record<string, unknown>;

  for (const word of expectedWords) {
    if (!(word in obj)) {
      throw new Error(`Missing key in LLM output: "${word}"`);
    }
    if (!Array.isArray(obj[word]) || (obj[word] as unknown[]).length === 0) {
      throw new Error(`LLM output for "${word}" is not a non-empty array`);
    }
  }

  return obj;
}

/**
 * Takes parsed LLM output and builds final enriched objects with composite IDs.
 */
export function buildEnrichedData(
  parsed: Record<string, unknown>,
  rawLanguage: string,
  rawPos: string,
): Map<string, EnrichedSense[]> {
  const language = (LANG_MAP[rawLanguage] || rawLanguage) as Language;
  const pos = (POS_MAP[rawPos] || rawPos) as Pos;
  const results = new Map<string, EnrichedSense[]>();

  for (const [word, sensesArray] of Object.entries(parsed)) {
    const senses = (sensesArray as unknown[]).map((item, index) => {
      const sense = item as Omit<
        EnrichedSense,
        "id" | "word" | "language" | "pos"
      >;

      return {
        id: `${word}:${language}:${pos}:${index}`,
        word,
        language,
        pos,
        ...sense,
      } as EnrichedSense;
    });

    results.set(word, senses);
  }

  return results;
}

/**
 * Enriches a batch of words by calling the LLM, parsing the response, and building final data.
 */
export async function enrichWord(
  words: string[],
  rawLanguage: string,
  rawPos: string,
): Promise<EnrichmentResult> {
  const llmResponse = await callLlm(words, rawLanguage, rawPos);
  const parsed = parseLlmResponse(llmResponse.content, words);
  const results = buildEnrichedData(parsed, rawLanguage, rawPos);

  return {
    results,
    metrics: {
      promptTokens: llmResponse.promptTokens,
      completionTokens: llmResponse.completionTokens,
      totalTokens: llmResponse.totalTokens,
      promptTimeMs: llmResponse.promptTimeMs,
      completionTimeMs: llmResponse.completionTimeMs,
    },
  };
}

/**
 * Enriches a batch of words with retry and split-on-failure logic.
 * Retries up to BATCH_CONFIG.maxRetries times, then splits batch in half and retries each half.
 * Continues splitting until batch size is 1, then throws if still failing.
 */
export async function enrichWordWithRetry(
  words: string[],
  rawLanguage: string,
  rawPos: string,
  attempt: number = 1,
): Promise<EnrichmentResult> {
  try {
    return await enrichWord(words, rawLanguage, rawPos);
  } catch (error: unknown) {
    const errorMessage = error instanceof Error ? error.message : String(error);

    if (words.length === 1) {
      throw new Error(
        `Failed to enrich word "${words[0]}" after ${attempt} attempts: ${errorMessage}`,
        { cause: error },
      );
    }

    if (attempt < BATCH_CONFIG.maxRetries) {
      console.log(
        `  Retry ${attempt}/${BATCH_CONFIG.maxRetries} for batch [${words.join(", ")}]: ${errorMessage}`,
      );
      return enrichWordWithRetry(words, rawLanguage, rawPos, attempt + 1);
    }

    // Max retries reached — split and retry
    console.log(
      `  Splitting batch [${words.join(", ")}] after ${BATCH_CONFIG.maxRetries} failed attempts`,
    );

    const half = Math.ceil(words.length / 2);
    const left = words.slice(0, half);
    const right = words.slice(half);

    const leftResult = await enrichWordWithRetry(left, rawLanguage, rawPos, 1);
    const rightResult = await enrichWordWithRetry(
      right,
      rawLanguage,
      rawPos,
      1,
    );

    // Merge results
    const merged = new Map([...leftResult.results, ...rightResult.results]);
    const mergedMetrics = {
      promptTokens:
        leftResult.metrics.promptTokens + rightResult.metrics.promptTokens,
      completionTokens:
        leftResult.metrics.completionTokens +
        rightResult.metrics.completionTokens,
      totalTokens:
        leftResult.metrics.totalTokens + rightResult.metrics.totalTokens,
      promptTimeMs:
        leftResult.metrics.promptTimeMs + rightResult.metrics.promptTimeMs,
      completionTimeMs:
        leftResult.metrics.completionTimeMs +
        rightResult.metrics.completionTimeMs,
    };

    return { results: merged, metrics: mergedMetrics };
  }
}