feat: progress bar and batch flag

Author: caiopizzol

apps/cli/commands/extract.ts

@@ -22,7 +22,7 @@ function parseFlags(args: string[]): ParsedFlags {
     const next = args[i + 1];
 
     switch (arg) {
-      case "--batch-size":
+      case "--batch":
       case "-b":
         flags.batchSize = parseInt(next || "", 10);
         i++;
@@ -55,7 +55,7 @@ Storage is auto-selected based on environment:
 Already-extracted files are automatically skipped (tracked in index.jsonl).
 
 Options
-  --batch-size, -b <n>    Number of files per batch (default: from EXTRACT_BATCH_SIZE or 100)
+  --batch, -b <n>         Limit to n documents (default: all)
   --workers, -w <n>       Number of parallel workers (default: from EXTRACT_WORKERS or 4)
   --verbose, -v           Show detailed progress
   --help, -h              Show this help
@@ -68,11 +68,11 @@ Environment Variables
   R2_BUCKET_NAME          R2 bucket (default: docx-corpus)
   EXTRACT_INPUT_PREFIX    Input prefix (default: documents)
   EXTRACT_OUTPUT_PREFIX   Output prefix (default: extracted)
-  EXTRACT_BATCH_SIZE      Batch size (default: 100)
   EXTRACT_WORKERS         Worker count (default: 4)
 
 Examples
-  corpus extract                    # Extract up to batch size
+  corpus extract                    # Extract all documents
+  corpus extract -b 100             # Limit to 100 documents
   corpus extract -v                 # With verbose output
   corpus extract -b 50 -w 8         # Custom batch/workers
 `;
@@ -101,7 +101,7 @@ export async function runExtract(args: string[]) {
     storage,
     inputPrefix: envConfig.extract.inputPrefix,
     outputPrefix: envConfig.extract.outputPrefix,
-    batchSize: flags.batchSize ?? envConfig.extract.batchSize,
+    batchSize: flags.batchSize ?? Infinity,
     workers: flags.workers ?? envConfig.extract.workers,
   };
 
@@ -113,7 +113,7 @@ export async function runExtract(args: string[]) {
   console.log(`Input:   ${config.inputPrefix}/`);
   console.log(`Output:  ${config.outputPrefix}/`);
   console.log(`Workers: ${config.workers}`);
-  console.log(`Batch:   ${config.batchSize}`);
+  console.log(`Batch:   ${config.batchSize === Infinity ? "all" : config.batchSize}`);
   console.log("");
 
   try {

packages/extractor/processor.ts

@@ -2,6 +2,7 @@ import { mkdir, rm } from "node:fs/promises";
 import { join, basename, dirname } from "node:path";
 import { tmpdir } from "node:os";
 import type { ExtractConfig, ExtractedDocument } from "./types";
+import { formatProgress, writeMultiLineProgress } from "@docx-corpus/shared";
 
 const PYTHON_DIR = join(dirname(import.meta.path), "python");
 const PYTHON_PATH = join(PYTHON_DIR, ".venv", "bin", "python");
@@ -34,7 +35,7 @@ export async function processDirectory(
       const id = extractIdFromKey(key);
       if (!indexState.successIds.has(id)) {
         files.push(key);
-        if (files.length >= batchSize) {
+        if (batchSize !== Infinity && files.length >= batchSize) {
           console.log(`Listed ${files.length} files to process (batch limit)`);
           break;
         }
@@ -122,6 +123,35 @@ async function processBatch(
   let errorCount = 0;
   const queue = [...keys];
 
+  // Progress tracking (only used when not verbose)
+  const startTime = Date.now();
+  let lastThroughputUpdate = startTime;
+  let docsAtLastUpdate = 0;
+  let currentDocsPerSec = 0;
+  let prevLineCount = 0;
+
+  const updateProgress = () => {
+    const now = Date.now();
+    const elapsed = (now - lastThroughputUpdate) / 1000;
+    if (elapsed >= 1) {
+      currentDocsPerSec = (successCount + errorCount - docsAtLastUpdate) / elapsed;
+      lastThroughputUpdate = now;
+      docsAtLastUpdate = successCount + errorCount;
+    }
+
+    const lines = formatProgress({
+      saved: successCount + errorCount,
+      total: keys.length,
+      docsPerSec: currentDocsPerSec,
+      failed: errorCount > 0 ? errorCount : undefined,
+      elapsedMs: now - startTime,
+    });
+
+    prevLineCount = writeMultiLineProgress(lines, prevLineCount);
+  };
+
+  const progressInterval = !verbose ? setInterval(updateProgress, 100) : null;
+
   const processFile = async (): Promise<void> => {
     while (queue.length > 0) {
       const sourceKey = queue.shift();
@@ -175,6 +205,14 @@ async function processBatch(
     .map(() => processFile());
 
   await Promise.all(workerPromises);
+
+  // Clean up progress display
+  if (progressInterval) {
+    clearInterval(progressInterval);
+    updateProgress(); // Final update
+    console.log(); // Move to next line
+  }
+
   return { successCount, errorCount };
 }
 

packages/shared/ui.ts

@@ -53,10 +53,10 @@ export interface ProgressStats {
   saved: number;
   total: number;
   docsPerSec: number;
-  currentRps: number;
-  skipped: number;
-  failed: number;
-  retried: number;
+  currentRps?: number;
+  skipped?: number;
+  failed?: number;
+  retried?: number;
   elapsedMs: number;
 }
 
@@ -76,10 +76,14 @@ export function formatProgress(stats: ProgressStats): string[] {
 
   // Line 2: Metrics
   const metrics: string[] = [];
-  metrics.push(`${docsPerSec.toFixed(1)}/s @ ${currentRps} RPS`);
-  if (skipped > 0) metrics.push(`${skipped} dup`);
-  if (failed > 0) metrics.push(`${failed} fail`);
-  if (retried > 0) metrics.push(`${retried} retried`);
+  if (currentRps !== undefined && currentRps > 0) {
+    metrics.push(`${docsPerSec.toFixed(1)}/s @ ${currentRps} RPS`);
+  } else {
+    metrics.push(`${docsPerSec.toFixed(1)}/s`);
+  }
+  if (skipped !== undefined && skipped > 0) metrics.push(`${skipped} dup`);
+  if (failed !== undefined && failed > 0) metrics.push(`${failed} fail`);
+  if (retried !== undefined && retried > 0) metrics.push(`${retried} retried`);
   metrics.push(formatDuration(elapsedMs));
 
   lines.push(metrics.join(" · "));