@ax-llm/ax 12.0.13 → 12.0.15

package/index.js

@@ -7536,6 +7536,22 @@ function parseFunctionCalls(ai, functionCalls, values, model) {
   }));
   return funcs;
 }
+function createFunctionConfig(functionList, definedFunctionCall, firstStep) {
+  let functionCall = definedFunctionCall;
+  if (!firstStep && (functionCall === "required" || typeof functionCall === "function")) {
+    return { functions: [], functionCall: void 0 };
+  }
+  if (!functionList) {
+    return { functions: [], functionCall };
+  }
+  const functions = functionList.map((f2) => {
+    if ("toFunction" in f2) {
+      return f2.toFunction();
+    }
+    return f2;
+  }).flat();
+  return { functions, functionCall };
+}
 
 // dsp/processResponse.ts
 import "stream/web";
@@ -9984,7 +10000,8 @@ var AxGen = class extends AxProgramWithSignature {
     mem,
     options,
     traceContext,
-    firstStep
+    functions,
+    functionCall
   }) {
     const {
       sessionId,
@@ -9992,8 +10009,6 @@ var AxGen = class extends AxProgramWithSignature {
       model,
       rateLimiter,
       stream,
-      functions: _functions,
-      functionCall: _functionCall,
       thinkingTokenBudget,
       showThoughts
     } = options ?? {};
@@ -10004,11 +10019,6 @@ var AxGen = class extends AxProgramWithSignature {
     if (chatPrompt.length === 0) {
       throw new Error("No chat prompt found");
     }
-    const functions = _functions?.map((f2) => "toFunction" in f2 ? f2.toFunction() : f2)?.flat();
-    let functionCall = _functionCall ?? this.options?.functionCall;
-    if (!firstStep && (functionCall === "required" || typeof functionCall === "function")) {
-      functionCall = void 0;
-    }
     const modelConfig = {
       ...options?.modelConfig,
       ...options?.sampleCount ? { n: options.sampleCount } : {},
@@ -10045,18 +10055,24 @@ var AxGen = class extends AxProgramWithSignature {
     span,
     traceContext
   }) {
-    const { sessionId, traceId, functions: _functions } = options ?? {};
+    const { sessionId, traceId, functions: functionList } = options ?? {};
+    const definedFunctionCall = options?.functionCall ?? this.options?.functionCall;
     const strictMode = options?.strictMode ?? false;
     const model = options.model;
     const states = this.createStates(options.sampleCount ?? 1);
     const usage = this.usage;
-    const functions = _functions?.map((f2) => "toFunction" in f2 ? f2.toFunction() : f2)?.flat() ?? [];
+    const { functions, functionCall } = createFunctionConfig(
+      functionList,
+      definedFunctionCall,
+      firstStep
+    );
     const res = await this.forwardSendRequest({
       ai,
       mem,
       options,
       traceContext,
-      firstStep
+      functions,
+      functionCall
     });
     if (res instanceof ReadableStream3) {
       yield* processStreamingResponse({
@@ -11700,13 +11716,6 @@ var AxBaseOptimizer = class {
     if (this.logger) {
       return this.logger;
     }
-    try {
-      const aiLogger = this.studentAI.getLogger();
-      if (aiLogger) {
-        return aiLogger;
-      }
-    } catch {
-    }
     return axDefaultOptimizerLogger;
   }
   /**
@@ -13398,6 +13407,11 @@ var AxMiPRO = class extends AxBaseOptimizer {
   bayesianOptimization;
   acquisitionFunction;
   explorationWeight;
+  // Self-consistency / multiple sampling
+  sampleCount;
+  // Surrogate model state for Bayesian optimization
+  miproConfigHistory = [];
+  surrogateModel = /* @__PURE__ */ new Map();
   constructor(args) {
     super(args);
     const options = args.options || {};
@@ -13419,6 +13433,7 @@ var AxMiPRO = class extends AxBaseOptimizer {
     this.bayesianOptimization = options.bayesianOptimization ?? false;
     this.acquisitionFunction = options.acquisitionFunction ?? "expected_improvement";
     this.explorationWeight = options.explorationWeight ?? 0.1;
+    this.sampleCount = options.sampleCount ?? 1;
     this.stats.convergenceInfo.convergenceThreshold = this.minImprovementThreshold;
   }
   /**
@@ -13463,43 +13478,186 @@ var AxMiPRO = class extends AxBaseOptimizer {
     ];
   }
   /**
-   * Generates instruction candidates using the teacher model if available
+   * Generates program summary for context-aware instruction generation
+   */
+  async generateProgramSummary(program, ai) {
+    let signature = "input -> output";
+    if ("getSignature" in program && typeof program.getSignature === "function") {
+      signature = program.getSignature();
+    }
+    const summaryPrompt = `
+Analyze this language model program and provide a concise summary of its purpose and structure.
+
+Program Signature: ${signature}
+
+Provide a 2-3 sentence summary focusing on:
+1. The main task or purpose of this program
+2. The input-output relationship
+3. Any special constraints or requirements
+
+Summary:`;
+    try {
+      const response = await ai.chat({
+        chatPrompt: [{ role: "user", content: summaryPrompt }]
+      });
+      if ("results" in response) {
+        return response.results[0]?.content?.trim() || "General language model program";
+      }
+      return "General language model program";
+    } catch {
+      return "General language model program";
+    }
+  }
+  /**
+   * Generates dataset summary for context-aware instruction generation
+   */
+  async generateDatasetSummary(examples, ai) {
+    if (examples.length === 0) return "No examples available";
+    const sampleSize = Math.min(this.viewDataBatchSize, examples.length);
+    const sampledExamples = examples.slice(0, sampleSize);
+    const exampleTexts = sampledExamples.map((ex, i) => `Example ${i + 1}: ${JSON.stringify(ex)}`).join("\n");
+    const summaryPrompt = `
+Analyze this dataset and provide a concise summary of its characteristics.
+
+Sample Examples:
+${exampleTexts}
+
+Provide a 2-3 sentence summary focusing on:
+1. The type of data and domain
+2. Common patterns or structures in the examples
+3. Key challenges or requirements for processing this data
+
+Dataset Summary:`;
+    try {
+      const response = await ai.chat({
+        chatPrompt: [{ role: "user", content: summaryPrompt }]
+      });
+      if ("results" in response) {
+        return response.results[0]?.content?.trim() || "General dataset";
+      }
+      return "General dataset";
+    } catch {
+      return "General dataset";
+    }
+  }
+  /**
+   * Enhanced instruction generation using AI with program and data awareness
+   */
+  async generateInstruction({
+    tip,
+    candidateIndex,
+    ai,
+    programSummary,
+    datasetSummary,
+    previousInstructions = []
+  }) {
+    let contextInfo = "";
+    if (this.programAwareProposer && programSummary) {
+      contextInfo += `
+Program Context: ${programSummary}`;
+    }
+    if (this.dataAwareProposer && datasetSummary) {
+      contextInfo += `
+Dataset Context: ${datasetSummary}`;
+    }
+    if (this.fewshotAwareProposer && previousInstructions.length > 0) {
+      contextInfo += `
+Previous Instructions (avoid repeating): ${previousInstructions.slice(-3).join("; ")}`;
+    }
+    const instructionPrompt = `
+Generate a high-quality instruction for a language model program.
+
+${contextInfo}
+
+${tip ? `Tip: ${tip}` : ""}
+
+Requirements:
+1. Be specific and actionable
+2. Focus on accuracy and clarity
+3. Consider the program's purpose and data characteristics
+4. Make the instruction distinct from previous ones
+5. Keep it concise but comprehensive
+
+Generate a single, well-crafted instruction:
+Instruction:`;
+    try {
+      const response = await ai.chat({
+        chatPrompt: [
+          {
+            role: "user",
+            content: instructionPrompt
+          }
+        ]
+      });
+      if ("results" in response) {
+        const instruction2 = response.results[0]?.content?.trim();
+        if (instruction2 && instruction2.length > 10) {
+          return instruction2;
+        }
+      }
+    } catch (error) {
+      if (this.isLoggingEnabled()) {
+        this.getLogger()?.(`Failed to generate AI instruction: ${error}`, {
+          tags: ["optimizer", "warning"]
+        });
+      }
+    }
+    const enhancedTemplates = [
+      "Analyze the input systematically and provide a precise, well-reasoned response.",
+      "Think through this step-by-step, considering all relevant factors before responding.",
+      "Examine the input carefully and generate an accurate, detailed answer.",
+      "Process the information methodically and deliver a clear, comprehensive response.",
+      "Consider the context thoroughly and provide a thoughtful, accurate answer."
+    ];
+    let instruction = enhancedTemplates[candidateIndex % enhancedTemplates.length] || enhancedTemplates[0];
+    if (tip) {
+      instruction = `${instruction} ${tip}`;
+    }
+    return instruction;
+  }
+  /**
+   * Generates instruction candidates using enhanced AI-powered generation
    * @param options Optional compile options that may override teacher AI
    * @returns Array of generated instruction candidates
    */
-  async proposeInstructionCandidates(options) {
+  async proposeInstructionCandidates(program, options) {
     const instructions = [];
     const aiToUse = this.getTeacherOrStudentAI(options);
+    let programSummary;
+    let datasetSummary;
+    if (this.programAwareProposer) {
+      programSummary = await this.generateProgramSummary(program, aiToUse);
+      if (this.isLoggingEnabled(options)) {
+        this.getLogger(options)?.(`Program summary: ${programSummary}`, {
+          tags: ["optimizer", "config"]
+        });
+      }
+    }
+    if (this.dataAwareProposer) {
+      datasetSummary = await this.generateDatasetSummary(this.examples, aiToUse);
+      if (this.isLoggingEnabled(options)) {
+        this.getLogger(options)?.(`Dataset summary: ${datasetSummary}`, {
+          tags: ["optimizer", "config"]
+        });
+      }
+    }
     const tips = this.tipAwareProposer ? this.generateTips() : [];
     for (let i = 0; i < this.numCandidates; i++) {
       const tipIndex = tips.length > 0 ? i % tips.length : -1;
-      const tipToUse = tipIndex >= 0 ? tips[tipIndex] : "";
+      const tipToUse = tipIndex >= 0 ? tips[tipIndex] : void 0;
       const instruction = await this.generateInstruction({
         tip: tipToUse,
         candidateIndex: i,
-        ai: aiToUse
+        ai: aiToUse,
+        programSummary,
+        datasetSummary,
+        previousInstructions: instructions
+        // Pass previous instructions for diversity
       });
       instructions.push(instruction);
     }
     return instructions;
   }
-  async generateInstruction({
-    tip,
-    candidateIndex
-  }) {
-    const baseInstructions = [
-      "Analyze the input carefully and provide a detailed response.",
-      "Think step by step and provide a clear answer.",
-      "Consider all aspects of the input before responding.",
-      "Provide a concise but comprehensive response.",
-      "Focus on accuracy and clarity in your response."
-    ];
-    let instruction = baseInstructions[candidateIndex % baseInstructions.length] || baseInstructions[0];
-    if (tip) {
-      instruction = `${instruction} ${tip}`;
-    }
-    return instruction;
-  }
   /**
    * Bootstraps few-shot examples for the program
    */
@@ -13544,7 +13702,7 @@ var AxMiPRO = class extends AxBaseOptimizer {
   /**
    * Runs optimization to find the best combination of few-shot examples and instructions
    */
-  async runOptimization(program, bootstrappedDemos, labeledExamples, instructions, valset, metricFn, options) {
+  async runOptimization(program, bootstrappedDemos, labeledExamples, instructions, validationExamples, metricFn, options) {
     let bestConfig = {
       instruction: instructions[0] || "",
       bootstrappedDemos: Math.min(1, bootstrappedDemos.length),
@@ -13580,25 +13738,37 @@ var AxMiPRO = class extends AxBaseOptimizer {
       );
     }
     for (let i = startRound; i < this.numTrials; i++) {
-      const config = {
-        instruction: instructions[i % instructions.length] || instructions[0] || "",
-        bootstrappedDemos: Math.min(
-          Math.floor(Math.random() * (bootstrappedDemos.length + 1)),
-          this.maxBootstrappedDemos
-        ),
-        labeledExamples: Math.min(
-          Math.floor(Math.random() * (labeledExamples.length + 1)),
-          this.maxLabeledDemos
-        )
-      };
+      let config;
+      if (this.bayesianOptimization && this.miproConfigHistory.length > 2) {
+        config = await this.selectConfigurationViaBayesianOptimization(
+          instructions,
+          bootstrappedDemos,
+          labeledExamples
+        );
+      } else {
+        config = {
+          instruction: instructions[i % instructions.length] || instructions[0] || "",
+          bootstrappedDemos: Math.min(
+            Math.floor(Math.random() * (bootstrappedDemos.length + 1)),
+            this.maxBootstrappedDemos
+          ),
+          labeledExamples: Math.min(
+            Math.floor(Math.random() * (labeledExamples.length + 1)),
+            this.maxLabeledDemos
+          )
+        };
+      }
       const score = await this.evaluateConfig(
         program,
         config,
         bootstrappedDemos,
         labeledExamples,
-        valset,
-        metricFn
+        validationExamples,
+        metricFn,
+        i + 1
+        // Pass current trial number for adaptive evaluation
       );
+      this.updateSurrogateModel(config, score);
       scoreHistory.push(score);
       const improvement = score - bestScore;
       if (improvement > this.minImprovementThreshold) {
@@ -13680,7 +13850,7 @@ var AxMiPRO = class extends AxBaseOptimizer {
     this.stats.convergenceInfo.converged = stagnationRounds < this.earlyStoppingTrials;
     return { bestConfig, bestScore };
   }
-  async evaluateConfig(program, config, bootstrappedDemos, labeledExamples, valset, metricFn) {
+  async evaluateConfig(program, config, bootstrappedDemos, labeledExamples, validationExamples, metricFn, currentTrial = 0) {
     const testProgram = { ...program };
     this.applyConfigToProgram(
       testProgram,
@@ -13690,12 +13860,31 @@ var AxMiPRO = class extends AxBaseOptimizer {
     );
     let totalScore = 0;
     let count = 0;
-    const evalSet = valset.slice(0, Math.min(5, valset.length));
+    let evalSize;
+    if (this.minibatch) {
+      const baseSize = Math.min(this.minibatchSize, validationExamples.length);
+      const isFullEvalTrial = currentTrial % this.minibatchFullEvalSteps === 0;
+      if (isFullEvalTrial || currentTrial > this.numTrials * 0.8) {
+        evalSize = Math.min(validationExamples.length, baseSize * 2);
+      } else {
+        evalSize = Math.max(3, Math.min(baseSize, validationExamples.length));
+      }
+    } else {
+      evalSize = validationExamples.length;
+    }
+    const evalIndices = this.shuffleArray([
+      ...Array(validationExamples.length).keys()
+    ]).slice(0, evalSize);
+    const evalSet = evalIndices.map((i) => validationExamples[i]);
     for (const example of evalSet) {
       try {
         const prediction = await testProgram.forward(
           this.studentAI,
-          example
+          example,
+          this.sampleCount > 1 ? {
+            sampleCount: this.sampleCount,
+            resultPicker: axMajorityVotePicker()
+          } : void 0
         );
         const score = await metricFn({ prediction, example });
         totalScore += score;
@@ -13707,6 +13896,17 @@ var AxMiPRO = class extends AxBaseOptimizer {
     }
     return count > 0 ? totalScore / count : 0;
   }
+  /**
+   * Fisher-Yates shuffle for stochastic evaluation
+   */
+  shuffleArray(array) {
+    const shuffled = [...array];
+    for (let i = shuffled.length - 1; i > 0; i--) {
+      const j = Math.floor(Math.random() * (i + 1));
+      [shuffled[i], shuffled[j]] = [shuffled[j], shuffled[i]];
+    }
+    return shuffled;
+  }
   applyConfigToProgram(program, config, bootstrappedDemos, labeledExamples) {
     if (program.setInstruction) {
       program.setInstruction(config.instruction);
@@ -13728,14 +13928,14 @@ var AxMiPRO = class extends AxBaseOptimizer {
     if (miproOptions?.auto) {
       this.configureAuto(miproOptions.auto);
     }
-    const valset = this.getValidationSet(options) || (miproOptions?.valset ?? this.examples.slice(0, Math.floor(this.examples.length * 0.2)));
+    const validationExamples = this.getValidationSet(options) || (miproOptions?.validationExamples ?? this.examples.slice(0, Math.floor(this.examples.length * 0.2)));
     if (this.isLoggingEnabled(options)) {
       this.getLogger(options)?.(
         `Starting MIPROv2 optimization with ${this.numTrials} trials`,
         { tags: ["optimizer", "start"] }
       );
       this.getLogger(options)?.(
-        `Using ${this.examples.length} examples for training and ${valset.length} for validation`,
+        `Using ${this.examples.length} examples for training and ${validationExamples.length} for validation`,
         { tags: ["optimizer", "config"] }
       );
       if (this.teacherAI) {
@@ -13765,7 +13965,10 @@ var AxMiPRO = class extends AxBaseOptimizer {
         );
       }
     }
-    const instructions = await this.proposeInstructionCandidates(options);
+    const instructions = await this.proposeInstructionCandidates(
+      program,
+      options
+    );
     if (this.isLoggingEnabled(options)) {
       this.getLogger(options)?.(
         `Generated ${instructions.length} instruction candidates`,
@@ -13783,7 +13986,7 @@ var AxMiPRO = class extends AxBaseOptimizer {
       bootstrappedDemos,
       labeledExamples,
       instructions,
-      valset,
+      validationExamples,
       metricFn,
       options
     );
@@ -13842,7 +14045,8 @@ var AxMiPRO = class extends AxBaseOptimizer {
         bootstrappedDemos: bestConfig.bootstrappedDemos,
         labeledExamples: bestConfig.labeledExamples,
         numCandidates: this.numCandidates,
-        numTrials: this.numTrials
+        numTrials: this.numTrials,
+        sampleCount: this.sampleCount
       }
     };
   }
@@ -13887,7 +14091,8 @@ var AxMiPRO = class extends AxBaseOptimizer {
       minImprovementThreshold: this.minImprovementThreshold,
       bayesianOptimization: this.bayesianOptimization,
       acquisitionFunction: this.acquisitionFunction,
-      explorationWeight: this.explorationWeight
+      explorationWeight: this.explorationWeight,
+      sampleCount: this.sampleCount
     };
   }
   /**
@@ -13922,12 +14127,17 @@ var AxMiPRO = class extends AxBaseOptimizer {
     if (config.minImprovementThreshold !== void 0) {
       this.minImprovementThreshold = config.minImprovementThreshold;
     }
+    if (config.sampleCount !== void 0) {
+      this.sampleCount = config.sampleCount;
+    }
   }
   /**
    * Reset optimizer state for reuse with different programs
    */
   reset() {
     super.reset();
+    this.miproConfigHistory = [];
+    this.surrogateModel.clear();
     this.stats.convergenceInfo.convergenceThreshold = this.minImprovementThreshold;
   }
   /**
@@ -13945,8 +14155,8 @@ var AxMiPRO = class extends AxBaseOptimizer {
         "Reduce maxBootstrappedDemos or maxLabeledDemos, or provide more examples"
       );
     }
-    const valSetSize = this.getValidationSet().length;
-    if (valSetSize < 5) {
+    const validationSetSize = this.getValidationSet().length;
+    if (validationSetSize < 5) {
       result.issues.push(
         "Validation set too small for reliable MiPRO optimization"
       );
@@ -13960,6 +14170,141 @@ var AxMiPRO = class extends AxBaseOptimizer {
       suggestions: result.suggestions
     };
   }
+  /**
+   * Encodes a configuration into a string key for surrogate model lookup
+   */
+  encodeConfiguration(config) {
+    return `${config.instruction.length}_${config.bootstrappedDemos}_${config.labeledExamples}`;
+  }
+  /**
+   * Updates the surrogate model with a new configuration-score pair
+   */
+  updateSurrogateModel(config, score) {
+    this.miproConfigHistory.push({ config: { ...config }, score });
+    const key = this.encodeConfiguration(config);
+    const similarConfigs = this.miproConfigHistory.filter(
+      (entry) => this.encodeConfiguration(entry.config) === key
+    );
+    if (similarConfigs.length > 0) {
+      const scores = similarConfigs.map((entry) => entry.score);
+      const mean = scores.reduce((sum, s2) => sum + s2, 0) / scores.length;
+      const variance = scores.length > 1 ? scores.reduce((sum, s2) => sum + Math.pow(s2 - mean, 2), 0) / (scores.length - 1) : 0.1;
+      this.surrogateModel.set(key, { mean, variance });
+    }
+  }
+  /**
+   * Predicts performance using the surrogate model
+   */
+  predictPerformance(config) {
+    const key = this.encodeConfiguration(config);
+    if (this.surrogateModel.has(key)) {
+      return this.surrogateModel.get(key);
+    }
+    if (this.miproConfigHistory.length > 0) {
+      const similarities = this.miproConfigHistory.map((entry) => {
+        const diff = Math.abs(entry.config.bootstrappedDemos - config.bootstrappedDemos) + Math.abs(entry.config.labeledExamples - config.labeledExamples);
+        return { score: entry.score, similarity: 1 / (1 + diff) };
+      });
+      const totalWeight = similarities.reduce((sum, s2) => sum + s2.similarity, 0);
+      const weightedMean = similarities.reduce((sum, s2) => sum + s2.score * s2.similarity, 0) / totalWeight;
+      return { mean: weightedMean, variance: 0.2 };
+    }
+    return { mean: 0.5, variance: 0.3 };
+  }
+  /**
+   * Calculates acquisition function value for Bayesian optimization
+   */
+  calculateAcquisitionValue(config) {
+    const prediction = this.predictPerformance(config);
+    const { mean, variance } = prediction;
+    const std = Math.sqrt(variance);
+    const bestScore = this.miproConfigHistory.length > 0 ? Math.max(...this.miproConfigHistory.map((entry) => entry.score)) : 0;
+    switch (this.acquisitionFunction) {
+      case "expected_improvement": {
+        const improvement = mean - bestScore;
+        if (std === 0) return Math.max(0, improvement);
+        const z = improvement / std;
+        const phi = 0.5 * (1 + this.erf(z / Math.sqrt(2)));
+        const pdfValue = Math.exp(-0.5 * z * z) / Math.sqrt(2 * Math.PI);
+        return improvement * phi + std * pdfValue;
+      }
+      case "upper_confidence_bound": {
+        return mean + this.explorationWeight * std;
+      }
+      case "probability_improvement": {
+        const improvement = mean - bestScore;
+        if (std === 0) return improvement > 0 ? 1 : 0;
+        const z = improvement / std;
+        return 0.5 * (1 + this.erf(z / Math.sqrt(2)));
+      }
+      default:
+        return mean;
+    }
+  }
+  /**
+   * Error function approximation for acquisition function calculations
+   */
+  erf(x) {
+    const a1 = 0.254829592;
+    const a2 = -0.284496736;
+    const a3 = 1.421413741;
+    const a4 = -1.453152027;
+    const a5 = 1.061405429;
+    const p = 0.3275911;
+    const sign = x >= 0 ? 1 : -1;
+    x = Math.abs(x);
+    const t = 1 / (1 + p * x);
+    const y = 1 - ((((a5 * t + a4) * t + a3) * t + a2) * t + a1) * t * Math.exp(-x * x);
+    return sign * y;
+  }
+  /**
+   * Selects the next configuration to evaluate using Bayesian optimization
+   */
+  async selectConfigurationViaBayesianOptimization(instructions, bootstrappedDemos, labeledExamples) {
+    const candidates = [];
+    const numCandidates = Math.min(20, instructions.length * 3);
+    for (let i = 0; i < numCandidates; i++) {
+      const config = {
+        instruction: instructions[i % instructions.length] || instructions[0] || "",
+        bootstrappedDemos: Math.min(
+          Math.floor(Math.random() * (bootstrappedDemos.length + 1)),
+          this.maxBootstrappedDemos
+        ),
+        labeledExamples: Math.min(
+          Math.floor(Math.random() * (labeledExamples.length + 1)),
+          this.maxLabeledDemos
+        )
+      };
+      const acquisitionValue = this.calculateAcquisitionValue(config);
+      candidates.push({ config, acquisitionValue });
+    }
+    candidates.sort((a, b) => b.acquisitionValue - a.acquisitionValue);
+    return candidates[0].config;
+  }
+};
+var axMajorityVotePicker = () => {
+  return async (data) => {
+    if (data.type === "fields") {
+      const counts = {};
+      for (const { index, sample } of data.results) {
+        const key = JSON.stringify(sample);
+        if (!counts[key]) {
+          counts[key] = { count: 0, index };
+        }
+        counts[key].count += 1;
+      }
+      let bestKey;
+      let bestCount = -1;
+      for (const [k, v] of Object.entries(counts)) {
+        if (v.count > bestCount) {
+          bestCount = v.count;
+          bestKey = k;
+        }
+      }
+      return counts[bestKey]?.index ?? 0;
+    }
+    return data.results[0]?.index ?? 0;
+  };
 };
 
 // ai/mock/api.ts