wip

data-pipeline/utils/enrich-word.ts

@@ -0,0 +1,240 @@
+// utils/enrich-word.ts
+import { ENRICHMENT_SYSTEM_PROMPT } 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";
+
+const LANG_MAP: Record<string, string> = {
+  english: "en",
+  italian: "it",
+  german: "de",
+  french: "fr",
+  spanish: "es",
+};
+
+const POS_MAP: Record<string, string> = {
+  nouns: "noun",
+  verbs: "verb",
+  adverbs: "adverb",
+  adjectives: "adjective",
+};
+
+interface LlmResponse {
+  content: string;
+  promptTokens: number;
+  completionTokens: number;
+  totalTokens: number;
+  promptTimeMs: number;
+  completionTimeMs: number;
+  promptTokensPerSecond: number;
+  completionTokensPerSecond: number;
+}
+
+export interface EnrichmentResult {
+  results: Map<string, EnrichedSense[]>;
+  metrics: {
+    promptTokens: number;
+    completionTokens: number;
+    totalTokens: number;
+    promptTimeMs: number;
+    completionTimeMs: number;
+    promptTokensPerSecond: number;
+    completionTokensPerSecond: number;
+  };
+}
+
+/**
+ * Calls the local LLM with the enrichment prompt.
+ * Returns the response content and timing metrics.
+ */
+async function callLlm(words: string[]): Promise<LlmResponse> {
+  const adapter = createAdapter();
+  return adapter.call(words, ENRICHMENT_SYSTEM_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);
+  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,
+      promptTokensPerSecond: llmResponse.promptTokensPerSecond,
+      completionTokensPerSecond: llmResponse.completionTokensPerSecond,
+    },
+  };
+}
+
+/**
+ * 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,
+      promptTokensPerSecond:
+        (leftResult.metrics.promptTokensPerSecond +
+          rightResult.metrics.promptTokensPerSecond) /
+        2,
+      completionTokensPerSecond:
+        (leftResult.metrics.completionTokensPerSecond +
+          rightResult.metrics.completionTokensPerSecond) /
+        2,
+    };
+
+    return { results: merged, metrics: mergedMetrics };
+  }
+}