@jaex/dstsx 0.1.0 → 0.1.2

package/dist/index.cjs

@@ -99,6 +99,7 @@ __export(index_exports, {
   evaluate: () => evaluate,
   exactMatch: () => exactMatch,
   f1: () => f1,
+  firstPrediction: () => firstPrediction,
   majority: () => majority,
   passAtK: () => passAtK,
   rouge: () => rouge,
@@ -490,7 +491,7 @@ var LM = class {
   #cache;
   #diskCache;
   #requestCount = 0;
-  #tokenUsage = { promptTokens: 0, completionTokens: 0, totalTokens: 0 };
+  #tokenUsage = { promptTokens: 0, completionTokens: 0, totalTokens: 0, cachedPromptTokens: 0 };
   constructor(model, cacheOptions = {}) {
     this.model = model;
     this.#cache = new LRUCache(cacheOptions.maxSize, cacheOptions.ttlMs);
@@ -524,6 +525,9 @@ var LM = class {
       this.#tokenUsage.promptTokens += response.usage.promptTokens;
       this.#tokenUsage.completionTokens += response.usage.completionTokens;
       this.#tokenUsage.totalTokens += response.usage.totalTokens;
+      if (response.usage.cachedPromptTokens) {
+        this.#tokenUsage.cachedPromptTokens += response.usage.cachedPromptTokens;
+      }
     }
     return response;
   }
@@ -570,21 +574,28 @@ var LM = class {
 // src/lm/adapters/OpenAI.ts
 var OpenAI = class extends LM {
   #options;
+  #client;
   constructor(options = {}) {
     super(options.model ?? "gpt-4o");
     this.#options = options;
   }
+  async #getClient() {
+    if (!this.#client) {
+      const { default: Client } = await import("openai").catch(() => {
+        throw new Error(
+          "The `openai` package is required for the OpenAI adapter.\nInstall it with: npm install openai"
+        );
+      });
+      this.#client = new Client({
+        apiKey: this.#options.apiKey ?? process.env["OPENAI_API_KEY"],
+        baseURL: this.#options.baseURL,
+        maxRetries: this.#options.maxRetries ?? 3
+      });
+    }
+    return this.#client;
+  }
   async _call(prompt, config) {
-    const { default: OpenAIClient } = await import("openai").catch(() => {
-      throw new Error(
-        "The `openai` package is required for the OpenAI adapter.\nInstall it with: npm install openai"
-      );
-    });
-    const client = new OpenAIClient({
-      apiKey: this.#options.apiKey ?? process.env["OPENAI_API_KEY"],
-      baseURL: this.#options.baseURL,
-      maxRetries: this.#options.maxRetries ?? 3
-    });
+    const client = await this.#getClient();
     const messages = typeof prompt === "string" ? [{ role: "user", content: prompt }] : prompt;
     const response = await client.chat.completions.create({
       model: config.model ?? this.model,
@@ -598,28 +609,21 @@ var OpenAI = class extends LM {
     const texts = (response.choices ?? []).map(
       (c) => c.message?.content ?? ""
     );
+    const usageDetails = response.usage;
     return {
       text: texts[0] ?? "",
       texts,
-      usage: response.usage ? {
-        promptTokens: response.usage.prompt_tokens,
-        completionTokens: response.usage.completion_tokens,
-        totalTokens: response.usage.total_tokens
+      usage: usageDetails ? {
+        promptTokens: usageDetails.prompt_tokens,
+        completionTokens: usageDetails.completion_tokens,
+        totalTokens: usageDetails.total_tokens,
+        ...usageDetails.prompt_tokens_details?.cached_tokens ? { cachedPromptTokens: usageDetails.prompt_tokens_details.cached_tokens } : {}
       } : null,
       raw: response
     };
   }
   async *stream(prompt, config = {}) {
-    const { default: OpenAIClient } = await import("openai").catch(() => {
-      throw new Error(
-        "The `openai` package is required for the OpenAI adapter.\nInstall it with: npm install openai"
-      );
-    });
-    const client = new OpenAIClient({
-      apiKey: this.#options.apiKey ?? process.env["OPENAI_API_KEY"],
-      baseURL: this.#options.baseURL,
-      maxRetries: this.#options.maxRetries ?? 3
-    });
+    const client = await this.#getClient();
     const messages = typeof prompt === "string" ? [{ role: "user", content: prompt }] : prompt;
     const stream = await client.chat.completions.create({
       model: config.model ?? this.model,
@@ -640,64 +644,164 @@ var OpenAI = class extends LM {
   }
 };
 
+// src/settings/Settings.ts
+var import_node_async_hooks = require("async_hooks");
+var contextStore = new import_node_async_hooks.AsyncLocalStorage();
+var Settings = class {
+  #global = {};
+  // ---------------------------------------------------------------------------
+  // Effective settings: async-context overrides take precedence over globals.
+  // ---------------------------------------------------------------------------
+  get #current() {
+    const ctx = contextStore.getStore();
+    return ctx !== void 0 ? { ...this.#global, ...ctx } : this.#global;
+  }
+  // ---------------------------------------------------------------------------
+  // Accessors
+  // ---------------------------------------------------------------------------
+  get lm() {
+    return this.#current.lm;
+  }
+  get rm() {
+    return this.#current.rm;
+  }
+  get lmConfig() {
+    return this.#current.lmConfig;
+  }
+  get logLevel() {
+    return this.#current.logLevel ?? "warn";
+  }
+  get cacheDir() {
+    return this.#current.cacheDir;
+  }
+  // ---------------------------------------------------------------------------
+  // Mutation
+  // ---------------------------------------------------------------------------
+  /**
+   * Merge `options` into the global settings.  Existing keys are overwritten;
+   * omitted keys are unchanged.  This does NOT affect currently running
+   * {@link Settings.context} scopes.
+   */
+  configure(options) {
+    this.#global = { ...this.#global, ...options };
+  }
+  /**
+   * Reset all global settings to their defaults.
+   */
+  reset() {
+    this.#global = {};
+  }
+  /**
+   * Return a deep-frozen snapshot of the currently effective settings
+   * (respects any active async-context overrides).
+   */
+  inspect() {
+    return Object.freeze({ ...this.#current });
+  }
+  /**
+   * Run `fn` inside an async-context-local settings scope.
+   *
+   * The `overrides` are merged on top of the current global settings and
+   * stored in an `AsyncLocalStorage` context.  Concurrent calls each get
+   * their own isolated snapshot — they never overwrite each other's settings.
+   *
+   * @example
+   * ```ts
+   * // In an Express/Fastify handler:
+   * await settings.context({ lm: perRequestLM }, () => program.forward(inputs));
+   * ```
+   */
+  async context(overrides, fn) {
+    const merged = { ...this.#global, ...overrides };
+    return contextStore.run(merged, fn);
+  }
+};
+var settings = new Settings();
+
 // src/lm/adapters/Anthropic.ts
 var Anthropic = class extends LM {
   #options;
+  #client;
   constructor(options = {}) {
     super(options.model ?? "claude-3-5-sonnet-20241022");
     this.#options = options;
   }
+  async #getClient() {
+    if (!this.#client) {
+      const { default: Client } = await import("@anthropic-ai/sdk").catch(() => {
+        throw new Error(
+          "The `@anthropic-ai/sdk` package is required for the Anthropic adapter.\nInstall it with: npm install @anthropic-ai/sdk"
+        );
+      });
+      this.#client = new Client({
+        apiKey: this.#options.apiKey ?? process.env["ANTHROPIC_API_KEY"],
+        maxRetries: this.#options.maxRetries ?? 3
+      });
+    }
+    return this.#client;
+  }
   async _call(prompt, config) {
-    const { default: Anthropic2 } = await import("@anthropic-ai/sdk").catch(() => {
-      throw new Error(
-        "The `@anthropic-ai/sdk` package is required for the Anthropic adapter.\nInstall it with: npm install @anthropic-ai/sdk"
-      );
-    });
-    const client = new Anthropic2({
-      apiKey: this.#options.apiKey ?? process.env["ANTHROPIC_API_KEY"],
-      maxRetries: this.#options.maxRetries ?? 3
-    });
+    const client = await this.#getClient();
+    const doCache = config.promptCaching ?? settings.lmConfig?.promptCaching ?? this.#options.promptCaching ?? false;
     const msgs = typeof prompt === "string" ? [{ role: "user", content: prompt }] : prompt;
     const systemMsg = msgs.find((m) => m.role === "system");
     const userMsgs = msgs.filter((m) => m.role !== "system");
+    let systemPayload;
+    if (systemMsg) {
+      systemPayload = doCache ? [{ type: "text", text: systemMsg.content, cache_control: { type: "ephemeral" } }] : systemMsg.content;
+    }
+    const messagesPayload = userMsgs.map((m, idx) => {
+      const isLast = idx === userMsgs.length - 1;
+      return {
+        role: m.role,
+        content: doCache && isLast ? [{ type: "text", text: m.content, cache_control: { type: "ephemeral" } }] : m.content
+      };
+    });
     const response = await client.messages.create({
       model: config.model ?? this.model,
       max_tokens: config.maxTokens ?? 1024,
-      system: systemMsg?.content,
-      messages: userMsgs.map((m) => ({ role: m.role, content: m.content })),
+      system: systemPayload,
+      messages: messagesPayload,
       temperature: config.temperature,
       ...config.extra ?? {}
     });
     const text = response.content.filter((b) => b.type === "text").map((b) => b.text ?? "").join("") ?? "";
+    const usage = response.usage;
+    const cachedInput = (usage.cache_read_input_tokens ?? 0) + (usage.cache_creation_input_tokens ?? 0);
     return {
       text,
       texts: [text],
-      usage: response.usage ? {
-        promptTokens: response.usage.input_tokens,
-        completionTokens: response.usage.output_tokens,
-        totalTokens: response.usage.input_tokens + response.usage.output_tokens
+      usage: usage ? {
+        promptTokens: usage.input_tokens,
+        completionTokens: usage.output_tokens,
+        totalTokens: usage.input_tokens + usage.output_tokens,
+        ...cachedInput > 0 ? { cachedPromptTokens: cachedInput } : {}
       } : null,
       raw: response
     };
   }
   async *stream(prompt, config = {}) {
-    const { default: Anthropic2 } = await import("@anthropic-ai/sdk").catch(() => {
-      throw new Error(
-        "The `@anthropic-ai/sdk` package is required for the Anthropic adapter.\nInstall it with: npm install @anthropic-ai/sdk"
-      );
-    });
-    const client = new Anthropic2({
-      apiKey: this.#options.apiKey ?? process.env["ANTHROPIC_API_KEY"],
-      maxRetries: this.#options.maxRetries ?? 3
-    });
+    const client = await this.#getClient();
+    const doCache = config.promptCaching ?? settings.lmConfig?.promptCaching ?? this.#options.promptCaching ?? false;
     const msgs = typeof prompt === "string" ? [{ role: "user", content: prompt }] : prompt;
     const systemMsg = msgs.find((m) => m.role === "system");
     const userMsgs = msgs.filter((m) => m.role !== "system");
+    let systemPayload;
+    if (systemMsg) {
+      systemPayload = doCache ? [{ type: "text", text: systemMsg.content, cache_control: { type: "ephemeral" } }] : systemMsg.content;
+    }
+    const messagesPayload = userMsgs.map((m, idx) => {
+      const isLast = idx === userMsgs.length - 1;
+      return {
+        role: m.role,
+        content: doCache && isLast ? [{ type: "text", text: m.content, cache_control: { type: "ephemeral" } }] : m.content
+      };
+    });
     const stream = client.messages.stream({
       model: config.model ?? this.model,
       max_tokens: config.maxTokens ?? 1024,
-      system: systemMsg?.content,
-      messages: userMsgs.map((m) => ({ role: m.role, content: m.content })),
+      system: systemPayload,
+      messages: messagesPayload,
       ...config.extra ?? {}
     });
     for await (const event of stream) {
@@ -951,6 +1055,9 @@ var MockLM = class extends LM {
 };
 
 // src/modules/Module.ts
+function firstPrediction(result) {
+  return Array.isArray(result) ? result[0] ?? new Prediction({}) : result;
+}
 var Module = class _Module {
   /**
    * Recursively discover all {@link Predict} sub-modules by walking the own
@@ -1014,80 +1121,6 @@ var Module = class _Module {
   }
 };
 
-// src/settings/Settings.ts
-var import_node_async_hooks = require("async_hooks");
-var contextStore = new import_node_async_hooks.AsyncLocalStorage();
-var Settings = class {
-  #global = {};
-  // ---------------------------------------------------------------------------
-  // Effective settings: async-context overrides take precedence over globals.
-  // ---------------------------------------------------------------------------
-  get #current() {
-    const ctx = contextStore.getStore();
-    return ctx !== void 0 ? { ...this.#global, ...ctx } : this.#global;
-  }
-  // ---------------------------------------------------------------------------
-  // Accessors
-  // ---------------------------------------------------------------------------
-  get lm() {
-    return this.#current.lm;
-  }
-  get rm() {
-    return this.#current.rm;
-  }
-  get lmConfig() {
-    return this.#current.lmConfig;
-  }
-  get logLevel() {
-    return this.#current.logLevel ?? "warn";
-  }
-  get cacheDir() {
-    return this.#current.cacheDir;
-  }
-  // ---------------------------------------------------------------------------
-  // Mutation
-  // ---------------------------------------------------------------------------
-  /**
-   * Merge `options` into the global settings.  Existing keys are overwritten;
-   * omitted keys are unchanged.  This does NOT affect currently running
-   * {@link Settings.context} scopes.
-   */
-  configure(options) {
-    this.#global = { ...this.#global, ...options };
-  }
-  /**
-   * Reset all global settings to their defaults.
-   */
-  reset() {
-    this.#global = {};
-  }
-  /**
-   * Return a deep-frozen snapshot of the currently effective settings
-   * (respects any active async-context overrides).
-   */
-  inspect() {
-    return Object.freeze({ ...this.#current });
-  }
-  /**
-   * Run `fn` inside an async-context-local settings scope.
-   *
-   * The `overrides` are merged on top of the current global settings and
-   * stored in an `AsyncLocalStorage` context.  Concurrent calls each get
-   * their own isolated snapshot — they never overwrite each other's settings.
-   *
-   * @example
-   * ```ts
-   * // In an Express/Fastify handler:
-   * await settings.context({ lm: perRequestLM }, () => program.forward(inputs));
-   * ```
-   */
-  async context(overrides, fn) {
-    const merged = { ...this.#global, ...overrides };
-    return contextStore.run(merged, fn);
-  }
-};
-var settings = new Settings();
-
 // src/modules/Predict.ts
 var Predict = class extends Module {
   signature;
@@ -1107,9 +1140,7 @@ var Predict = class extends Module {
   async forward(inputs) {
     const lm = settings.lm;
     if (!lm) {
-      throw new Error(
-        "No LM configured. Call settings.configure({ lm }) before using Predict."
-      );
+      throw new Error("No LM configured. Call settings.configure({ lm }) before using Predict.");
     }
     const prompt = this.#buildPrompt(inputs);
     const config = settings.lmConfig ?? {};
@@ -1124,7 +1155,8 @@ var Predict = class extends Module {
    */
   async *stream(inputs) {
     const lm = settings.lm;
-    if (!lm) throw new Error("No LM configured. Call settings.configure({ lm }) before using Predict.");
+    if (!lm)
+      throw new Error("No LM configured. Call settings.configure({ lm }) before using Predict.");
     const prompt = this.#buildPrompt(inputs);
     const config = settings.lmConfig ?? {};
     yield* lm.stream(prompt, config);
@@ -1141,9 +1173,7 @@ var Predict = class extends Module {
   }
   load(state) {
     if (Array.isArray(state["demos"])) {
-      this.demos = state["demos"].map(
-        (d) => new Example(d)
-      );
+      this.demos = state["demos"].map((d) => new Example(d));
     }
     if (typeof state["instructions"] === "string") {
       this.instructions = state["instructions"];
@@ -1164,32 +1194,80 @@ var Predict = class extends Module {
       }
     }
     lines.push(this.#formatInputs(inputs));
-    for (const [name] of this.signature.outputs) {
-      lines.push(`${name}:`);
+    for (const [name, meta] of this.signature.outputs) {
+      const label = meta.prefix ?? `${name}:`;
+      lines.push(label.endsWith(":") ? label : `${label}:`);
     }
     return lines.join("\n");
   }
   #formatExample(data) {
-    return [...this.signature.inputs, ...this.signature.outputs].map(([name]) => `${name}: ${String(data[name] ?? "")}`).join("\n");
+    const formatField = (name, meta) => {
+      const label = meta?.prefix ?? `${name}:`;
+      const prefix = label.endsWith(":") ? label : `${label}:`;
+      return `${prefix} ${String(data[name] ?? "")}`;
+    };
+    return [
+      ...[...this.signature.inputs].map(([name, meta]) => formatField(name, meta)),
+      ...[...this.signature.outputs].map(([name, meta]) => formatField(name, meta))
+    ].join("\n");
   }
   #formatInputs(inputs) {
-    return [...this.signature.inputs].map(([name]) => `${name}: ${String(inputs[name] ?? "")}`).join("\n");
+    return [...this.signature.inputs].map(([name, meta]) => {
+      const label = meta.prefix ?? `${name}:`;
+      const prefix = label.endsWith(":") ? label : `${label}:`;
+      return `${prefix} ${String(inputs[name] ?? "")}`;
+    }).join("\n");
   }
   /**
    * Parse a raw completion string into a map of field name → value.
    *
-   * Looks for `fieldName: <value>` lines in the completion.
+   * Looks for `fieldName: <value>` or `prefix: <value>` patterns in the completion.
+   * Handles common LLM formatting like **bold**, *italic*, etc.
+   * Supports multi-line values by extracting content between field markers.
    */
   #parseCompletion(text) {
     const result = {};
-    const outputKeys = [...this.signature.outputs.keys()];
-    for (const key of outputKeys) {
-      const regex = new RegExp(`^${key}:\\s*(.*)$`, "mi");
-      const match = regex.exec(text);
-      if (match) {
-        result[key] = (match[1] ?? "").trim();
+    const outputs = [...this.signature.outputs.entries()];
+    const fieldLabels = [];
+    for (const [key, meta] of outputs) {
+      const prefixBase = (meta.prefix ?? "").replace(/:$/, "").trim();
+      const patterns = [key];
+      if (prefixBase && prefixBase.toLowerCase() !== key.toLowerCase()) {
+        patterns.push(prefixBase);
+      }
+      fieldLabels.push({ key, patterns });
+    }
+    const allPatterns = fieldLabels.flatMap((f) => f.patterns);
+    const allPatternsEscaped = allPatterns.map((p) => p.replace(/[.*+?^${}()|[\]\\]/g, "\\$&"));
+    const markerPattern = `(?:\\*\\*|\\*|_)?(${allPatternsEscaped.join("|")})(?:\\*\\*|\\*|_)?:`;
+    const markerRegex = new RegExp(markerPattern, "gi");
+    const markers = [];
+    let markerMatch;
+    while ((markerMatch = markerRegex.exec(text)) !== null) {
+      const matchedLabel = markerMatch[1].toLowerCase();
+      for (const { key, patterns } of fieldLabels) {
+        if (patterns.some((p) => p.toLowerCase() === matchedLabel)) {
+          markers.push({
+            key,
+            start: markerMatch.index,
+            end: markerMatch.index + markerMatch[0].length
+          });
+          break;
+        }
       }
     }
+    for (let i = 0; i < markers.length; i++) {
+      const marker = markers[i];
+      const contentStart = marker.end;
+      const contentEnd = i + 1 < markers.length ? markers[i + 1].start : text.length;
+      let content = text.slice(contentStart, contentEnd).trim();
+      content = content.replace(/^(\*\*|\*|_)+\s*/, "");
+      content = content.replace(/\s*(\*\*|\*|_)+$/, "");
+      if (!(marker.key in result)) {
+        result[marker.key] = content;
+      }
+    }
+    const outputKeys = outputs.map(([k]) => k);
     if (outputKeys.length === 1 && !(outputKeys[0] in result)) {
       result[outputKeys[0]] = text.trim();
     }
@@ -1235,9 +1313,6 @@ var ChainOfThoughtWithHint = class extends ChainOfThought {
     });
     Object.assign(this, { signature: extendedSig });
   }
-  async forward(inputs) {
-    return super.forward(inputs);
-  }
 };
 
 // src/modules/MultiChainComparison.ts
@@ -1255,14 +1330,23 @@ var MultiChainComparison = class extends Module {
   }
   async forward(inputs) {
     const completions = [];
+    const outputKey = [...this.#cot.signature.outputs.keys()].find((k) => k !== "rationale") ?? "answer";
     for (let i = 0; i < this.M; i++) {
       const result = await this.#cot.forward(inputs);
-      const outputKey = [...this.#cot.signature.outputs.keys()].find((k) => k !== "rationale");
-      completions.push(String(result.get(outputKey ?? "answer") ?? ""));
+      completions.push(String(result.get(outputKey) ?? ""));
     }
-    return this.#aggregator.forward({
+    const aggregated = await this.#aggregator.forward({
       completions: completions.map((c, i) => `Option ${i + 1}: ${c}`).join("\n")
     });
+    const rawAnswer = String(aggregated.get("answer") ?? "");
+    const optionMatch = /^Option\s*(\d+)/i.exec(rawAnswer);
+    if (optionMatch) {
+      const optionIdx = parseInt(optionMatch[1], 10) - 1;
+      if (optionIdx >= 0 && optionIdx < completions.length) {
+        return new Prediction({ answer: completions[optionIdx] ?? rawAnswer });
+      }
+    }
+    return new Prediction({ answer: rawAnswer });
   }
 };
 
@@ -1276,7 +1360,6 @@ var ReAct = class extends Module {
     this.tools = new Map(tools.map((t) => [t.name, t]));
     this.maxIter = maxIter;
     const toolDescriptions = tools.map((t) => `${t.name}: ${t.description}`).join("\n");
-    const base = typeof signature === "string" ? signature : signature;
     const instructions = `You are an agent. Use the following tools:
 ${toolDescriptions}
 
@@ -1286,9 +1369,12 @@ Action: <tool>[<args>]
 Observation: <result>
 ...
 Finish[<answer>]`;
-    this.#predictor = new Predict(
-      typeof base === "string" ? `${base}` : base
-    );
+    let baseSig = typeof signature === "string" ? Signature.from(signature) : signature;
+    baseSig = baseSig.withInput("trajectory", {
+      description: "Interleaved reasoning steps so far",
+      optional: true
+    });
+    this.#predictor = new Predict(baseSig);
     this.#predictor.instructions = instructions;
   }
   async forward(inputs) {
@@ -1371,10 +1457,11 @@ var ProgramOfThought = class extends Module {
     for (let attempt = 0; attempt < this.maxAttempts; attempt++) {
       const genInputs = attempt === 0 ? {
         ...inputs,
-        instructions: "Write JavaScript code to compute the answer. Use a `return` statement for the final value."
+        instructions: "Write JavaScript code to compute the answer. IMPORTANT: Use a `return` statement (not console.log) for the final value. Do not include markdown code fences."
       } : { code, error: lastError };
       const generated = attempt === 0 ? await this.#codeGenerator.forward(genInputs) : await this.#corrector.forward(genInputs);
       code = String(generated.get("code") ?? generated.get("fixed_code") ?? "");
+      code = code.replace(/^```(?:\w+)?\s*\n?/i, "").replace(/\n?```\s*$/i, "").trim();
       try {
         if (this.sandbox === "worker") {
           result = await this.#executeInWorker(code, this.timeoutMs);
@@ -1513,7 +1600,9 @@ var Retry = class extends Module {
     let lastError;
     for (let attempt = 0; attempt < this.maxAttempts; attempt++) {
       try {
-        return await this.#inner.forward(...args);
+        return firstPrediction(
+          await this.#inner.forward(...args)
+        );
       } catch (err) {
         if (err instanceof AssertionError) {
           lastError = err;
@@ -1548,7 +1637,7 @@ var BestOfN = class extends Module {
         () => this.#inner.forward(...args)
       )
     );
-    return this.#reduce(results);
+    return this.#reduce(results.map(firstPrediction));
   }
 };
 
@@ -1567,7 +1656,7 @@ var Ensemble = class extends Module {
         (m) => m.forward(...args)
       )
     );
-    return this.#reduce(results);
+    return this.#reduce(results.map(firstPrediction));
   }
 };
 
@@ -1676,6 +1765,7 @@ var TypedChainOfThought = class extends TypedPredictor {
 };
 
 // src/modules/Parallel.ts
+var import_node_timers = require("timers");
 var Parallel = class extends Module {
   #modules;
   #timeoutMs;
@@ -1684,29 +1774,37 @@ var Parallel = class extends Module {
     this.#modules = modules;
     this.#timeoutMs = options.timeoutMs;
   }
-  /** Run all modules in parallel and return all predictions. */
+  /**
+   * Run all modules in parallel and return one {@link Prediction} per module.
+   * If a module's `forward()` returns multiple predictions, the first is used.
+   */
   async run(...args) {
     const tasks = this.#modules.map(
       (m) => m.forward(...args)
     );
+    let settled;
     if (this.#timeoutMs !== void 0) {
       const timeoutMs = this.#timeoutMs;
       const withTimeout = tasks.map(
         (t) => Promise.race([
           t,
           new Promise(
-            (_, reject) => setTimeout(() => reject(new Error("Parallel: timeout")), timeoutMs)
+            (_, reject) => (0, import_node_timers.setTimeout)(() => reject(new Error("Parallel: timeout")), timeoutMs)
           )
         ])
       );
-      return Promise.all(withTimeout);
+      settled = Promise.all(withTimeout);
+    } else {
+      settled = Promise.all(tasks);
     }
-    return Promise.all(tasks);
+    return (await settled).map(firstPrediction);
   }
-  /** For Module interface compatibility — returns first prediction. */
+  /**
+   * Execute all modules in parallel and return all predictions as an array
+   * (one entry per module).
+   */
   async forward(...args) {
-    const results = await this.run(...args);
-    return results[0];
+    return this.run(...args);
   }
 };
 
@@ -1727,7 +1825,7 @@ var Refine = class extends Module {
   }
   async forward(...args) {
     const innerForward = this.#inner.forward.bind(this.#inner);
-    let prediction = await innerForward(...args);
+    let prediction = firstPrediction(await innerForward(...args));
     for (let i = 0; i < this.#maxRefinements; i++) {
       if (this.#stopCondition?.(prediction)) break;
       const outputStr = JSON.stringify(prediction.toDict());
@@ -1750,7 +1848,7 @@ var Refine = class extends Module {
         };
       }
       try {
-        prediction = await innerForward(...newArgs);
+        prediction = firstPrediction(await innerForward(...newArgs));
       } catch {
         break;
       }
@@ -2875,6 +2973,7 @@ var JsonFileTracker = class extends Tracker {
   evaluate,
   exactMatch,
   f1,
+  firstPrediction,
   majority,
   passAtK,
   rouge,