lila/data-pipeline/pipeline.ts

import { isWordProcessed } from "./utils/check-if-json-exists.js";
import { createBaseJson } from "./utils/create-base-json.js";
import { ensureOutputFolders } from "./utils/create-output-dirs.js";
import { scanSourceData } from "./utils/scanning-source-files.js";
import { createLineReader } from "./utils/create-line-reader.js";
import { checkLlmServer } from "./utils/check-llm-server.js";
import { getWordFilePath } from "./utils/get-word-file-path.js";
import { mergeEnrichedData } from "./utils/merge-enriched-data.js";
import { enrichWordWithRetry } from "./utils/enrich-word.js";
import { writeJsonFile } from "./utils/write-json-file.js";
import { deleteFileIfExists } from "./utils/delete-file.js";
import { PipelineTimer } from "./utils/pipeline-timer.js";
import { ProgressTracker } from "./utils/progress-tracker.js";
import { verifyEnrichedFile } from "./utils/verify-enriched-file.js";
import { BATCH_CONFIG } from "./config/batch.js";

async function main() {
  console.log("Starting data pipeline...\n");

  const timer = new PipelineTimer();

  // step 1: scanning for source files
  console.log("\n step 1: scanning the source files...");
  const wordlists = scanSourceData(import.meta.dirname);

  // step 2: ensuring output folders exist
  console.log("\n step 2: creating necessary output folders...");
  ensureOutputFolders(wordlists);

  // step 3: check to verify the local AI engine is ready before touching anything
  console.log("\n step 3: verifying local AI engine status...");
  await checkLlmServer();

  // Step 4: Loop through the wordlists array
  console.log("\n step 4: looping through the wordlists...");

  for (const wordlist of wordlists) {
    console.log(
      `\nReading list: [${wordlist.language.toUpperCase()}] -> [${wordlist.pos.toUpperCase()}]`,
    );

    const rl = createLineReader(wordlist.sourcePath);

    // Collect words and count them
    const words: string[] = [];
    for await (const line of rl) {
      const word = line.trim().toLowerCase();
      if (word) words.push(word);
    }

    // Filter out already-processed words
    const unprocessedWords = words.filter(
      (word) => !isWordProcessed(word, wordlist.outputDir),
    );

    const skippedCount = words.length - unprocessedWords.length;
    if (skippedCount > 0) {
      console.log(`  Skipped ${skippedCount} already-processed words`);
    }

    const progress = new ProgressTracker(unprocessedWords.length);

    // Step 5: Process in batches
    for (let i = 0; i < unprocessedWords.length; i += BATCH_CONFIG.size) {
      const batch = unprocessedWords.slice(i, i + BATCH_CONFIG.size);
      const batchNum = Math.floor(i / BATCH_CONFIG.size) + 1;
      const totalBatches = Math.ceil(
        unprocessedWords.length / BATCH_CONFIG.size,
      );
      const batchLabel = `Batch ${batchNum}/${totalBatches}`;

      console.log(`\n  ${batchLabel}: [${batch.join(", ")}]`);

      // Create skeletons for all words in batch
      for (const word of batch) {
        createBaseJson(
          word,
          wordlist.outputDir,
          wordlist.language,
          wordlist.pos,
        );
      }

      timer.startWord();

      try {
        // Step 6: enrich batch with senses (with retry/split)
        const result = await enrichWordWithRetry(
          batch,
          wordlist.language,
          wordlist.pos,
        );

        // Step 7: write each word's result
        for (const [word, senses] of result.results) {
          progress.next();
          console.log(
            `  ${progress.format(`Enriched and saved: ${word}.json`)}`,
          );

          const targetFilePath = getWordFilePath(word, wordlist.outputDir);
          const enrichedData = mergeEnrichedData(word, senses);

          writeJsonFile(targetFilePath, enrichedData);

          // Verify the generated file
          const verification = verifyEnrichedFile(targetFilePath);
          if (!verification.valid) {
            console.error(`  Warning: Schema violations in ${word}.json:`);
            for (const error of verification.errors) {
              console.error(`    - ${error}`);
            }
          }

          timer.recordProcessed({
            promptTokens: result.metrics.promptTokens,
            completionTokens: result.metrics.completionTokens,
            totalTokens: result.metrics.totalTokens,
            promptTimeMs: result.metrics.promptTimeMs,
            completionTimeMs: result.metrics.completionTimeMs,
          });
        }

        console.log(`  ${timer.getWordTiming()}`);
        // Show ETA every 5 batches or on the last batch
        if (batchNum % 5 === 0 || batchNum === totalBatches) {
          console.log(`  📊 ${timer.getEta(unprocessedWords.length)}`);
        }
      } catch (error: unknown) {
        const errorMessage =
          error instanceof Error ? error.message : String(error);
        console.error(
          `  Failed to enrich batch [${batch.join(", ")}]: ${errorMessage}`,
        );

        // Cleanup: delete skeleton files for failed batch
        for (const word of batch) {
          const targetFilePath = getWordFilePath(word, wordlist.outputDir);
          deleteFileIfExists(targetFilePath);
          console.log(`  Removed incomplete file: ${word}.json`);
          progress.recordFailed();
        }

        timer.recordFailed();
      }
    }
  }

  timer.stop();
  console.log("\n" + timer.getSummary());
  console.log("\nGlobal data pipeline run completed successfully.");
}

// Fire the orchestrator block
main().catch((err) => {
  console.error("Critical unexpected pipeline failure:", err);
});