@ax-llm/ax 12.0.12 → 12.0.14

package/index.cjs

@@ -6257,15 +6257,24 @@ var MemoryImpl = class {
       debugRequest(items, this.options?.debugHideSystemPrompt);
     }
   }
+  addFunctionResults(results) {
+    const chat = results.map(({ index, ...value }) => ({
+      index,
+      value: structuredClone(value)
+    }));
+    const lastItem = this.getLast();
+    if (lastItem?.role === "function") {
+      lastItem.chat.push(...chat);
+    } else {
+      this.data.push({ role: "function", chat });
+    }
+  }
   addResponse(results) {
-    const chat = results.map((result) => ({
-      index: result.index,
-      value: structuredClone(result)
+    const chat = results.map(({ index, ...value }) => ({
+      index,
+      value: structuredClone(value)
     }));
-    this.data.push({
-      role: "assistant",
-      chat
-    });
+    this.data.push({ role: "assistant", chat });
     if (this.options?.debug) {
       for (const result of results) {
         debugResponse(result);
@@ -6353,9 +6362,20 @@ var MemoryImpl = class {
   history(index) {
     const result = [];
     for (const { role, chat } of this.data) {
-      const value = chat.find((v) => v.index === index)?.value;
-      if (value) {
-        result.push({ role, ...value });
+      let values;
+      if (role === "function") {
+        values = chat.filter((v) => v.index === index).map((v) => v.value);
+      } else {
+        values = chat.find((v) => v.index === index)?.value;
+      }
+      if (Array.isArray(values)) {
+        result.push(
+          ...values.map(
+            (v) => ({ ...v, role })
+          )
+        );
+      } else if (values) {
+        result.push({ ...values, role });
       }
     }
     return result;
@@ -6393,20 +6413,8 @@ var AxMemory = class {
     axValidateChatResponseResult(results);
     this.getMemory(sessionId).addResponse(results);
   }
-  addFunctionResult({
-    functionId,
-    isError,
-    index,
-    result
-  }, sessionId) {
-    const functionMessage = {
-      role: "function",
-      functionId,
-      isError,
-      result
-    };
-    axValidateChatRequestMessage(functionMessage);
-    this.getMemory(sessionId).addRequest([functionMessage], index);
+  addFunctionResults(results, sessionId) {
+    this.getMemory(sessionId).addFunctionResults(results);
   }
   updateResult(result, sessionId) {
     this.getMemory(sessionId).updateResult(result);
@@ -7643,51 +7651,48 @@ var processFunctions = async ({
       }
       return {
         result: functionResult ?? "",
+        role: "function",
         functionId: func.id,
         index
       };
     }).catch((e) => {
-      if (e instanceof FunctionError) {
-        const result = e.getFixingInstructions();
-        if (span) {
-          const errorEventData = {
-            name: func.name,
-            message: e.toString()
-          };
-          if (!excludeContentFromTrace) {
-            errorEventData.args = func.args;
-            errorEventData.fixing_instructions = result;
-          }
-          span.addEvent("function.error", errorEventData);
+      if (!(e instanceof FunctionError)) {
+        throw e;
+      }
+      const result = e.getFixingInstructions();
+      if (span) {
+        const errorEventData = {
+          name: func.name,
+          message: e.toString()
+        };
+        if (!excludeContentFromTrace) {
+          errorEventData.args = func.args;
+          errorEventData.fixing_instructions = result;
         }
-        mem.addFunctionResult(
-          {
-            functionId: func.id,
-            isError: true,
-            index,
-            result
-          },
-          sessionId
-        );
-        mem.addTag("error", sessionId);
-        if (ai.getOptions().debug) {
-          const logger = ai.getLogger();
-          logger(`\u274C Function Error Correction:
+        span.addEvent("function.error", errorEventData);
+      }
+      if (ai.getOptions().debug) {
+        const logger = ai.getLogger();
+        logger(`\u274C Function Error Correction:
 ${result}`, {
-            tags: ["error"]
-          });
-        }
-      } else {
-        throw e;
+          tags: ["error"]
+        });
       }
+      return {
+        functionId: func.id,
+        isError: true,
+        index,
+        result,
+        role: "function"
+      };
     });
     return promise;
   });
   const results = await Promise.all(promises);
-  for (const result of results) {
-    if (result) {
-      mem.addFunctionResult(result, sessionId);
-    }
+  const functionResults = results.filter((result) => result !== void 0);
+  mem.addFunctionResults(functionResults, sessionId);
+  if (functionResults.some((result) => result.isError)) {
+    mem.addTag("error", sessionId);
   }
   return functionsExecuted;
 };
@@ -9961,6 +9966,96 @@ var AxProgram = class {
   }
 };
 
+// dsp/samples.ts
+function checkForFunctionCalls(mem, sessionId) {
+  const history = mem.history(0, sessionId);
+  const hasFunctionResults = history.some((msg) => msg.role === "function");
+  const hasFunctionCalls = history.some(
+    (msg) => msg.role === "assistant" && "functionCalls" in msg && Array.isArray(msg.functionCalls) && msg.functionCalls.length > 0
+  );
+  return hasFunctionCalls && hasFunctionResults;
+}
+function extractFunctionResults(mem, sessionId) {
+  const history = mem.history(0, sessionId);
+  const results = [];
+  const assistantMessages = history.filter(
+    (msg) => msg.role === "assistant" && "functionCalls" in msg && Array.isArray(msg.functionCalls) && msg.functionCalls.length > 0
+  );
+  const functionMessages = history.filter((msg) => msg.role === "function");
+  for (const assistantMsg of assistantMessages) {
+    if ("functionCalls" in assistantMsg && assistantMsg.functionCalls) {
+      for (const funcCall of assistantMsg.functionCalls) {
+        const funcResult = functionMessages.find(
+          (msg) => "functionId" in msg && msg.functionId === funcCall.id
+        );
+        if (funcResult && "result" in funcResult && "functionId" in funcResult) {
+          results.push({
+            index: results.length,
+            // Use sequential index for function results
+            functionName: funcCall.function.name,
+            functionId: funcCall.id,
+            args: funcCall.function.params || "",
+            result: String(funcResult.result),
+            isError: "isError" in funcResult ? Boolean(funcResult.isError) : false
+          });
+        }
+      }
+    }
+  }
+  return results;
+}
+async function selectFromSamples(buffer, options, mem, sessionId) {
+  if (!options?.resultPicker || buffer.length <= 1) {
+    return 0;
+  }
+  const resultPicker = options.resultPicker;
+  const hasFunctionCalls = mem ? checkForFunctionCalls(mem, sessionId) : false;
+  if (hasFunctionCalls && mem) {
+    const functionResults = extractFunctionResults(mem, sessionId);
+    const selectedIndex = await resultPicker({
+      type: "function",
+      results: functionResults
+    });
+    if (selectedIndex < 0 || selectedIndex >= functionResults.length) {
+      throw new Error(
+        `Result picker returned invalid index: ${selectedIndex}. Must be between 0 and ${functionResults.length - 1}`
+      );
+    }
+    return selectedIndex;
+  } else {
+    const fieldResults = buffer.map((b, index) => ({
+      index,
+      sample: b.delta
+    }));
+    const selectedIndex = await resultPicker({
+      type: "fields",
+      results: fieldResults
+    });
+    if (selectedIndex < 0 || selectedIndex >= buffer.length) {
+      throw new Error(
+        `Result picker returned invalid index: ${selectedIndex}. Must be between 0 and ${buffer.length - 1}`
+      );
+    }
+    return selectedIndex;
+  }
+}
+async function selectFromSamplesInMemory(mem, sessionId, options) {
+  const lastMemory = mem?.getLast(sessionId);
+  if (!lastMemory || lastMemory.role !== "assistant") {
+    return 0;
+  }
+  if (lastMemory.chat.length <= 1) {
+    return 0;
+  }
+  const buffer = lastMemory.chat.map((chat) => ({
+    version: 0,
+    index: chat.index,
+    delta: chat.value
+  }));
+  const selectedIndex = await selectFromSamples(buffer, options, mem, sessionId);
+  return selectedIndex;
+}
+
 // dsp/validate.ts
 function handleValidationError(mem, errorFields, ai, promptTemplate, sessionId) {
   mem.addRequest(
@@ -10076,7 +10171,10 @@ var AxGen = class extends AxProgramWithSignature {
       thinkingTokenBudget,
       showThoughts
     } = options ?? {};
-    const chatPrompt = mem?.history(0, sessionId) ?? [];
+    const selectedIndex = await selectFromSamplesInMemory(mem, sessionId, {
+      resultPicker: options?.resultPicker
+    });
+    const chatPrompt = mem?.history(selectedIndex, sessionId) ?? [];
     if (chatPrompt.length === 0) {
       throw new Error("No chat prompt found");
     }
@@ -10087,7 +10185,8 @@ var AxGen = class extends AxProgramWithSignature {
     }
     const modelConfig = {
       ...options?.modelConfig,
-      ...options?.sampleCount ? { n: options.sampleCount } : {}
+      ...options?.sampleCount ? { n: options.sampleCount } : {},
+      ...options?.sampleCount && options?.modelConfig?.temperature == 1 ? { temperature: 0.8 } : {}
     };
     const res = await ai.chat(
       {
@@ -10364,15 +10463,58 @@ var AxGen = class extends AxProgramWithSignature {
       currentVersion = delta.version;
       buffer = mergeDeltas(buffer, delta);
     }
-    const result = buffer[0]?.delta ?? {};
+    const selectedIndex = await selectFromSamples(
+      buffer,
+      {
+        resultPicker: options?.resultPicker
+      },
+      // Pass memory to enable function result selection
+      options?.mem,
+      options?.sessionId
+    );
+    const selectedResult = buffer[selectedIndex];
+    const result = selectedResult?.delta ?? {};
     this.trace = { ...values, ...result };
     return result;
   }
   async *streamingForward(ai, values, options) {
-    yield* this._forward1(ai, values, {
+    if (!options?.resultPicker) {
+      yield* this._forward1(ai, values, {
+        ...options,
+        stream: true
+      });
+      return;
+    }
+    const generator = this._forward1(ai, values, {
       ...options,
       stream: true
     });
+    let buffer = [];
+    let currentVersion = 0;
+    for await (const delta of generator) {
+      if (delta.version !== currentVersion) {
+        buffer = [];
+      }
+      currentVersion = delta.version;
+      buffer = mergeDeltas(buffer, delta);
+    }
+    const selectedIndex = await selectFromSamples(
+      buffer,
+      {
+        resultPicker: options?.resultPicker
+      },
+      // Pass memory to enable function result selection
+      options?.mem,
+      options?.sessionId
+    );
+    const selectedResult = buffer[selectedIndex];
+    if (selectedResult) {
+      yield {
+        version: currentVersion,
+        index: selectedIndex,
+        delta: selectedResult.delta
+      };
+    }
   }
   setExamples(examples, options) {
     super.setExamples(examples, options);
@@ -11732,13 +11874,6 @@ var AxBaseOptimizer = class {
     if (this.logger) {
       return this.logger;
     }
-    try {
-      const aiLogger = this.studentAI.getLogger();
-      if (aiLogger) {
-        return aiLogger;
-      }
-    } catch {
-    }
     return axDefaultOptimizerLogger;
   }
   /**
@@ -13430,6 +13565,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 || {};
@@ -13451,6 +13591,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;
   }
   /**
@@ -13495,43 +13636,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
    */
@@ -13576,7 +13860,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),
@@ -13612,25 +13896,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) {
@@ -13712,7 +14008,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,
@@ -13722,12 +14018,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;
@@ -13739,6 +14054,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);
@@ -13760,14 +14086,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) {
@@ -13797,7 +14123,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`,
@@ -13815,7 +14144,7 @@ var AxMiPRO = class extends AxBaseOptimizer {
       bootstrappedDemos,
       labeledExamples,
       instructions,
-      valset,
+      validationExamples,
       metricFn,
       options
     );
@@ -13874,7 +14203,8 @@ var AxMiPRO = class extends AxBaseOptimizer {
         bootstrappedDemos: bestConfig.bootstrappedDemos,
         labeledExamples: bestConfig.labeledExamples,
         numCandidates: this.numCandidates,
-        numTrials: this.numTrials
+        numTrials: this.numTrials,
+        sampleCount: this.sampleCount
       }
     };
   }
@@ -13919,7 +14249,8 @@ var AxMiPRO = class extends AxBaseOptimizer {
       minImprovementThreshold: this.minImprovementThreshold,
       bayesianOptimization: this.bayesianOptimization,
       acquisitionFunction: this.acquisitionFunction,
-      explorationWeight: this.explorationWeight
+      explorationWeight: this.explorationWeight,
+      sampleCount: this.sampleCount
     };
   }
   /**
@@ -13954,12 +14285,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;
   }
   /**
@@ -13977,8 +14313,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"
       );
@@ -13992,6 +14328,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