@ax-llm/ax 11.0.27 → 11.0.29

package/index.js

@@ -58,7 +58,7 @@ var AxSpanKindValues = /* @__PURE__ */ ((AxSpanKindValues2) => {
 
 // util/apicall.ts
 import {
-  ReadableStream,
+  ReadableStream as ReadableStream2,
   TextDecoderStream as TextDecoderStreamNative,
   TransformStream as TransformStream3
 } from "stream/web";
@@ -307,7 +307,7 @@ var apiCall = async (api, json) => {
   const metrics = createRequestMetrics();
   let timeoutId;
   const baseUrl = new URL(process.env["PROXY"] ?? api.url);
-  const apiPath = [baseUrl.pathname, api.name].filter(Boolean).join("/").replace(/\/+/g, "/");
+  const apiPath = `${[baseUrl.pathname, api.name].filter(Boolean).join("/").replace(/\/+/g, "/")}${baseUrl.search}`;
   const apiUrl = new URL(apiPath, baseUrl);
   const requestId = crypto.randomUUID();
   if (api.validateRequest) {
@@ -427,7 +427,7 @@ var apiCall = async (api, json) => {
         }
       });
       let closed = false;
-      return new ReadableStream({
+      return new ReadableStream2({
         start(controller2) {
           const reader = res.body.pipeThrough(new textDecoderStream()).pipeThrough(new SSEParser()).pipeThrough(trackingStream).getReader();
           async function read() {
@@ -3733,7 +3733,7 @@ var AxAI = class {
 };
 
 // dsp/generate.ts
-import { ReadableStream as ReadableStream2 } from "node:stream/web";
+import { ReadableStream as ReadableStream3 } from "node:stream/web";
 import { SpanKind as SpanKind2 } from "@opentelemetry/api";
 
 // ai/util.ts
@@ -4637,16 +4637,19 @@ var validateValue = (field, value) => {
 };
 function mergeProgramUsage(usages) {
   const usageMap = {};
-  usages.forEach((usage) => {
+  for (const usage of usages) {
     const key = `${usage.ai}:${usage.model}`;
     if (!usageMap[key]) {
       usageMap[key] = { ...usage };
-      return;
+      continue;
     }
-    usageMap[key].promptTokens += usage.promptTokens;
-    usageMap[key].completionTokens += usage.completionTokens;
-    usageMap[key].totalTokens += usage.totalTokens;
-  });
+    const currentUsage = usageMap[key];
+    if (currentUsage) {
+      currentUsage.promptTokens += usage.promptTokens;
+      currentUsage.completionTokens += usage.completionTokens;
+      currentUsage.totalTokens += usage.totalTokens;
+    }
+  }
   return Object.values(usageMap);
 }
 var parseMarkdownList = (input) => {
@@ -4667,7 +4670,6 @@ var parseMarkdownList = (input) => {
     } else if (numberedListRegex.test(trimmedLine)) {
       list.push(trimmedLine.replace(numberedListRegex, "").trim());
     } else if (list.length === 0) {
-      continue;
     } else {
       throw new Error("Could not parse markdown list: mixed content detected");
     }
@@ -4740,12 +4742,72 @@ function matchesContent(content, prefix, startIndex = 0, prefixCache = globalPre
   );
   for (let i = 0; i < prefixes.length - 1; i++) {
     const partialPrefix = prefixes[i];
-    if (contentEnd.endsWith(partialPrefix)) {
+    if (partialPrefix && contentEnd.endsWith(partialPrefix)) {
       return -2;
     }
   }
   return -1;
 }
+var formatTime = (ms) => {
+  const seconds = Math.floor(ms / 1e3);
+  if (seconds < 60) return `${seconds}s`;
+  const minutes = Math.floor(seconds / 60);
+  const remainingSeconds = seconds % 60;
+  if (minutes < 60) return `${minutes}m ${remainingSeconds}s`;
+  const hours = Math.floor(minutes / 60);
+  const remainingMinutes = minutes % 60;
+  return `${hours}h ${remainingMinutes}m ${remainingSeconds}s`;
+};
+var calculateETA = (current, total, elapsedMs) => {
+  if (current === 0) return "calculating...";
+  const msPerItem = elapsedMs / current;
+  const remainingItems = total - current;
+  const etaMs = msPerItem * remainingItems;
+  return formatTime(etaMs);
+};
+var updateDetailedProgress = (roundIndex, current, total, elapsedTime, example, stats, configInfo, result, error) => {
+  process.stdout.write("\r\x1B[K");
+  const percentage = (current / total * 100).toFixed(1);
+  const formattedTime = formatTime(elapsedTime);
+  const itemsPerSecond = elapsedTime > 0 ? (current / elapsedTime * 1e3).toFixed(2) : "0.00";
+  const eta = calculateETA(current, total, elapsedTime);
+  let output = `Round ${roundIndex + 1}/${configInfo.maxRounds}: ${current}/${total} (${percentage}%) [${formattedTime}, ${itemsPerSecond} it/s, ETA: ${eta}]`;
+  const successRate = stats.totalCalls > 0 ? stats.successfulDemos / stats.totalCalls * 100 : 0;
+  output += ` | Success: ${stats.successfulDemos}/${stats.totalCalls} (${successRate.toFixed(1)}%)`;
+  if (configInfo.verboseMode || configInfo.debugMode) {
+    if (configInfo.costMonitoring) {
+      output += `
+  Tokens: ~${stats.estimatedTokenUsage.toLocaleString()} total`;
+    }
+    output += `
+  Batch: ${Math.floor(current / configInfo.batchSize) + 1}/${Math.ceil(total / configInfo.batchSize)}`;
+    if (configInfo.earlyStoppingPatience > 0 && stats.earlyStopping) {
+      output += `
+  Best round: ${stats.earlyStopping.bestScoreRound + 1}, Patience: ${configInfo.earlyStoppingPatience}`;
+    }
+  }
+  if (configInfo.debugMode) {
+    const exampleKeys = Object.keys(example).map((k) => {
+      const valueStr = JSON.stringify(example[k]);
+      const truncated = valueStr.length > 30 ? `${valueStr.substring(0, 30)}...` : valueStr;
+      return `${k}: ${truncated}`;
+    }).join(", ");
+    output += `
+  Example: {${exampleKeys}}`;
+    if (error) {
+      output += `
+  ERROR: ${error.message}`;
+    } else if (result) {
+      const resultStr = JSON.stringify(result);
+      const truncatedResult = resultStr.length > 50 ? `${resultStr.substring(0, 50)}...` : resultStr;
+      output += `
+  Result: ${truncatedResult}`;
+    }
+    output += `
+  Temperature: ${(0.7 + 1e-3 * current).toFixed(3)}`;
+  }
+  console.log(output);
+};
 
 // dsp/program.ts
 var AxProgramWithSignature = class {
@@ -4987,7 +5049,7 @@ ${outputFields}`);
     demos
   }) => {
     const renderedExamples = examples ? [
-      { type: "text", text: "## Examples:\n" },
+      { type: "text", text: "\n\n## Examples\n" },
       ...this.renderExamples(examples)
     ] : [];
     const renderedDemos = demos ? this.renderDemos(demos) : [];
@@ -5070,6 +5132,9 @@ ${outputFields}`);
         );
       }
       const renderedItem = [...renderedInputItem, ...renderedOutputItem];
+      if (index > 0 && renderedItem.length > 0 && renderedItem[0]?.type === "text") {
+        list.push({ type: "text", text: "---\n\n" });
+      }
       renderedItem.forEach((v) => {
         if ("text" in v) {
           v.text = v.text + "\n";
@@ -5253,10 +5318,7 @@ var processValue = (field, value) => {
   if (typeof value === "string") {
     return value;
   }
-  if (Array.isArray(value)) {
-    return value;
-  }
-  return JSON.stringify(value);
+  return JSON.stringify(value, null, 2);
 };
 var toFieldType = (type) => {
   const baseType = (() => {
@@ -6062,7 +6124,7 @@ var AxGen = class extends AxProgramWithSignature {
       mem,
       options
     });
-    if (res instanceof ReadableStream2) {
+    if (res instanceof ReadableStream3) {
       yield* this.processStreamingResponse({
         ai,
         model,
@@ -6139,6 +6201,9 @@ var AxGen = class extends AxProgramWithSignature {
           content,
           streamingValidation
         );
+        if (skip) {
+          continue;
+        }
         if (this.streamingAsserts.length !== 0) {
           await assertStreamingAssertions(
             this.streamingAsserts,
@@ -6157,9 +6222,6 @@ var AxGen = class extends AxProgramWithSignature {
           );
         }
         yield* streamValues(this.signature, content, values, xstate);
-        if (skip) {
-          continue;
-        }
         await assertAssertions(this.asserts, values);
       }
       if (result.finishReason === "length") {
@@ -6997,60 +7059,147 @@ function validateModels2(services) {
 // dsp/optimize.ts
 var AxBootstrapFewShot = class {
   ai;
+  teacherAI;
   program;
   examples;
   maxRounds;
   maxDemos;
   maxExamples;
+  batchSize;
+  earlyStoppingPatience;
+  costMonitoring;
+  maxTokensPerGeneration;
+  verboseMode;
+  debugMode;
   traces = [];
+  stats = {
+    totalCalls: 0,
+    successfulDemos: 0,
+    estimatedTokenUsage: 0,
+    earlyStopped: false
+  };
   constructor({
     ai,
     program,
     examples = [],
     options
   }) {
-    if (examples.length == 0) {
+    if (examples.length === 0) {
       throw new Error("No examples found");
     }
     this.maxRounds = options?.maxRounds ?? 3;
     this.maxDemos = options?.maxDemos ?? 4;
     this.maxExamples = options?.maxExamples ?? 16;
+    this.batchSize = options?.batchSize ?? 1;
+    this.earlyStoppingPatience = options?.earlyStoppingPatience ?? 0;
+    this.costMonitoring = options?.costMonitoring ?? false;
+    this.maxTokensPerGeneration = options?.maxTokensPerGeneration ?? 0;
+    this.verboseMode = options?.verboseMode ?? true;
+    this.debugMode = options?.debugMode ?? false;
     this.ai = ai;
+    this.teacherAI = options?.teacherAI;
     this.program = program;
     this.examples = examples;
   }
   async compileRound(roundIndex, metricFn, options) {
     const st = (/* @__PURE__ */ new Date()).getTime();
     const maxDemos = options?.maxDemos ?? this.maxDemos;
-    const aiOpt = { modelConfig: { temperature: 0.7 } };
+    const aiOpt = {
+      modelConfig: {
+        temperature: 0.7
+      }
+    };
+    if (this.maxTokensPerGeneration > 0) {
+      aiOpt.modelConfig.max_tokens = this.maxTokensPerGeneration;
+    }
     const examples = randomSample(this.examples, this.maxExamples);
-    for (let i = 0; i < examples.length; i++) {
+    const previousSuccessCount = this.traces.length;
+    for (let i = 0; i < examples.length; i += this.batchSize) {
       if (i > 0) {
         aiOpt.modelConfig.temperature = 0.7 + 1e-3 * i;
       }
-      const ex = examples[i];
-      if (!ex) {
-        throw new Error("Invalid example");
-      }
-      const exList = [...examples.slice(0, i), ...examples.slice(i + 1)];
-      this.program.setExamples(exList);
-      const res = await this.program.forward(this.ai, ex, aiOpt);
-      const success = metricFn({ prediction: res, example: ex });
-      if (success) {
-        this.traces = [...this.traces, ...this.program.getTraces()];
+      const batch = examples.slice(i, i + this.batchSize);
+      for (const ex of batch) {
+        if (!ex) {
+          continue;
+        }
+        const exList = examples.filter((e) => e !== ex);
+        this.program.setExamples(exList);
+        const aiService = this.teacherAI || this.ai;
+        this.stats.totalCalls++;
+        let res;
+        let error;
+        try {
+          res = await this.program.forward(aiService, ex, aiOpt);
+          if (this.costMonitoring) {
+            this.stats.estimatedTokenUsage += JSON.stringify(ex).length / 4 + JSON.stringify(res).length / 4;
+          }
+          const success = metricFn({ prediction: res, example: ex });
+          if (success) {
+            this.traces = [...this.traces, ...this.program.getTraces()];
+            this.stats.successfulDemos++;
+          }
+        } catch (err) {
+          error = err;
+          res = {};
+        }
+        const current = i + examples.length * roundIndex + (batch.indexOf(ex) + 1);
+        const total = examples.length * this.maxRounds;
+        const et = (/* @__PURE__ */ new Date()).getTime() - st;
+        if (this.verboseMode || this.debugMode) {
+          const configInfo = {
+            maxRounds: this.maxRounds,
+            batchSize: this.batchSize,
+            earlyStoppingPatience: this.earlyStoppingPatience,
+            costMonitoring: this.costMonitoring,
+            verboseMode: this.verboseMode,
+            debugMode: this.debugMode
+          };
+          updateDetailedProgress(
+            roundIndex,
+            current,
+            total,
+            et,
+            ex,
+            this.stats,
+            configInfo,
+            res,
+            error
+          );
+        } else {
+          updateProgressBar(
+            current,
+            total,
+            this.traces.length,
+            et,
+            "Tuning Prompt",
+            30
+          );
+        }
+        if (this.traces.length >= maxDemos) {
+          return;
+        }
       }
-      const current = i + examples.length * roundIndex;
-      const total = examples.length * this.maxRounds;
-      const et = (/* @__PURE__ */ new Date()).getTime() - st;
-      updateProgressBar(
-        current,
-        total,
-        this.traces.length,
-        et,
-        "Tuning Prompt",
-        30
-      );
-      if (this.traces.length > maxDemos) {
+    }
+    if (this.earlyStoppingPatience > 0) {
+      const newSuccessCount = this.traces.length;
+      const improvement = newSuccessCount - previousSuccessCount;
+      if (!this.stats.earlyStopping) {
+        this.stats.earlyStopping = {
+          bestScoreRound: improvement > 0 ? roundIndex : 0,
+          patienceExhausted: false
+        };
+      } else if (improvement > 0) {
+        this.stats.earlyStopping.bestScoreRound = roundIndex;
+      } else if (roundIndex - this.stats.earlyStopping.bestScoreRound >= this.earlyStoppingPatience) {
+        this.stats.earlyStopping.patienceExhausted = true;
+        this.stats.earlyStopped = true;
+        if (this.verboseMode || this.debugMode) {
+          console.log(
+            `
+Early stopping triggered after ${roundIndex + 1} rounds. No improvement for ${this.earlyStoppingPatience} rounds.`
+          );
+        }
         return;
       }
     }
@@ -7058,8 +7207,17 @@ var AxBootstrapFewShot = class {
   async compile(metricFn, options) {
     const maxRounds = options?.maxRounds ?? this.maxRounds;
     this.traces = [];
+    this.stats = {
+      totalCalls: 0,
+      successfulDemos: 0,
+      estimatedTokenUsage: 0,
+      earlyStopped: false
+    };
     for (let i = 0; i < maxRounds; i++) {
       await this.compileRound(i, metricFn, options);
+      if (this.stats.earlyStopped) {
+        break;
+      }
     }
     if (this.traces.length === 0) {
       throw new Error(
@@ -7067,22 +7225,32 @@ var AxBootstrapFewShot = class {
       );
     }
     const demos = groupTracesByKeys(this.traces);
-    return demos;
+    return {
+      demos,
+      stats: this.stats
+    };
+  }
+  // Get optimization statistics
+  getStats() {
+    return this.stats;
   }
 };
 function groupTracesByKeys(programTraces) {
   const groupedTraces = /* @__PURE__ */ new Map();
   for (const programTrace of programTraces) {
     if (groupedTraces.has(programTrace.programId)) {
-      groupedTraces.get(programTrace.programId).push(programTrace.trace);
+      const traces = groupedTraces.get(programTrace.programId);
+      if (traces) {
+        traces.push(programTrace.trace);
+      }
     } else {
       groupedTraces.set(programTrace.programId, [programTrace.trace]);
     }
   }
   const programDemosArray = [];
-  groupedTraces.forEach((traces, programId) => {
+  for (const [programId, traces] of groupedTraces.entries()) {
     programDemosArray.push({ traces, programId });
-  });
+  }
   return programDemosArray;
 }
 var randomSample = (array, n) => {
@@ -8202,6 +8370,639 @@ var AxJSInterpreter = class {
   }
 };
 
+// dsp/mipro.ts
+var AxMiPRO = class {
+  ai;
+  program;
+  examples;
+  maxBootstrappedDemos;
+  maxLabeledDemos;
+  numCandidates;
+  initTemperature;
+  numTrials;
+  minibatch;
+  minibatchSize;
+  minibatchFullEvalSteps;
+  programAwareProposer;
+  dataAwareProposer;
+  viewDataBatchSize;
+  tipAwareProposer;
+  fewshotAwareProposer;
+  seed;
+  verbose;
+  bootstrapper;
+  earlyStoppingTrials;
+  minImprovementThreshold;
+  constructor({
+    ai,
+    program,
+    examples = [],
+    options
+  }) {
+    if (examples.length === 0) {
+      throw new Error("No examples found");
+    }
+    const miproOptions = options || {};
+    this.numCandidates = miproOptions.numCandidates ?? 5;
+    this.initTemperature = miproOptions.initTemperature ?? 0.7;
+    this.maxBootstrappedDemos = miproOptions.maxBootstrappedDemos ?? 3;
+    this.maxLabeledDemos = miproOptions.maxLabeledDemos ?? 4;
+    this.numTrials = miproOptions.numTrials ?? 30;
+    this.minibatch = miproOptions.minibatch ?? true;
+    this.minibatchSize = miproOptions.minibatchSize ?? 25;
+    this.minibatchFullEvalSteps = miproOptions.minibatchFullEvalSteps ?? 10;
+    this.programAwareProposer = miproOptions.programAwareProposer ?? true;
+    this.dataAwareProposer = miproOptions.dataAwareProposer ?? true;
+    this.viewDataBatchSize = miproOptions.viewDataBatchSize ?? 10;
+    this.tipAwareProposer = miproOptions.tipAwareProposer ?? true;
+    this.fewshotAwareProposer = miproOptions.fewshotAwareProposer ?? true;
+    this.seed = miproOptions.seed;
+    this.verbose = miproOptions.verbose ?? false;
+    this.earlyStoppingTrials = miproOptions.earlyStoppingTrials ?? 5;
+    this.minImprovementThreshold = miproOptions.minImprovementThreshold ?? 0.01;
+    this.ai = ai;
+    this.program = program;
+    this.examples = examples;
+    this.bootstrapper = new AxBootstrapFewShot({
+      ai,
+      program,
+      examples,
+      options: {
+        maxDemos: this.maxBootstrappedDemos,
+        maxRounds: 3,
+        // Default, or adjust based on your needs
+        verboseMode: this.verbose
+      }
+    });
+  }
+  /**
+   * Configures the optimizer for light, medium, or heavy optimization
+   * @param level The optimization level: "light", "medium", or "heavy"
+   */
+  configureAuto(level) {
+    switch (level) {
+      case "light":
+        this.numCandidates = 3;
+        this.numTrials = 10;
+        this.minibatch = true;
+        this.minibatchSize = 20;
+        break;
+      case "medium":
+        this.numCandidates = 5;
+        this.numTrials = 20;
+        this.minibatch = true;
+        this.minibatchSize = 25;
+        break;
+      case "heavy":
+        this.numCandidates = 7;
+        this.numTrials = 30;
+        this.minibatch = true;
+        this.minibatchSize = 30;
+        break;
+    }
+  }
+  /**
+   * Generates creative tips for instruction generation
+   */
+  generateTips() {
+    return [
+      "Be very specific and detailed in your instructions.",
+      "Focus on step-by-step reasoning in your instructions.",
+      "Provide clear constraints and guidelines in your instructions.",
+      "Keep your instructions concise and to the point.",
+      "Emphasize accuracy and precision in your instructions.",
+      "Include examples of good outputs in your instructions.",
+      "Focus on handling edge cases in your instructions.",
+      "Explicitly outline the reasoning process in your instructions."
+    ];
+  }
+  /**
+   * Generates instruction candidates for each predictor in the program
+   * @returns Array of generated instruction candidates
+   */
+  async proposeInstructionCandidates() {
+    const instructions = [];
+    let programContext = "";
+    if (this.programAwareProposer) {
+      programContext = await this.generateProgramSummary();
+    }
+    let dataContext = "";
+    if (this.dataAwareProposer) {
+      dataContext = await this.generateDataSummary();
+    }
+    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 instruction = await this.generateInstruction({
+        programContext,
+        dataContext,
+        tip: tipToUse,
+        candidateIndex: i
+      });
+      instructions.push(instruction);
+    }
+    return instructions;
+  }
+  /**
+   * Generates a summary of the program structure for instruction proposal
+   */
+  async generateProgramSummary() {
+    const prompt = `Summarize the following program structure. Focus on the signatures, 
+      input/output fields, and the purpose of each component. Identify key components 
+      that might benefit from better instructions.`;
+    const programStr = JSON.stringify(this.program);
+    const response = await this.ai.chat({
+      chatPrompt: [
+        { role: "system", content: prompt },
+        { role: "user", content: programStr }
+      ],
+      modelConfig: { temperature: 0.2 }
+    });
+    if (response instanceof ReadableStream) {
+      return "";
+    }
+    return response.results[0]?.content || "";
+  }
+  /**
+   * Generates a summary of the dataset for instruction proposal
+   */
+  async generateDataSummary() {
+    const sampleSize = Math.min(this.viewDataBatchSize, this.examples.length);
+    const sample = this.examples.slice(0, sampleSize);
+    const prompt = `Analyze the following dataset examples and provide a summary 
+      of key patterns, input-output relationships, and any specific challenges 
+      the data presents. Focus on what makes a good answer and what patterns should
+      be followed.`;
+    const dataStr = JSON.stringify(sample);
+    const response = await this.ai.chat({
+      chatPrompt: [
+        { role: "system", content: prompt },
+        { role: "user", content: dataStr }
+      ],
+      modelConfig: { temperature: 0.2 }
+    });
+    if (response instanceof ReadableStream) {
+      return "";
+    }
+    return response.results[0]?.content || "";
+  }
+  /**
+   * Generates a specific instruction candidate
+   */
+  async generateInstruction({
+    programContext,
+    dataContext,
+    tip,
+    candidateIndex
+  }) {
+    const prompt = `Create a high-quality instruction for an AI model performing the task described below.
+    
+    ${programContext ? `PROGRAM CONTEXT:
+${programContext}
+
+` : ""}
+    ${dataContext ? `DATA CONTEXT:
+${dataContext}
+
+` : ""}
+    ${tip ? `STYLE TIP: ${tip}
+
+` : ""}
+    
+    Your task is to craft a clear, effective instruction that will help the AI model generate
+    accurate outputs for this task. Instruction #${candidateIndex + 1}/${this.numCandidates}.
+    
+    The instruction should be detailed enough to guide the model but not overly prescriptive
+    or restrictive. Focus on what makes a good response rather than listing exact steps.
+    
+    INSTRUCTION:`;
+    const response = await this.ai.chat({
+      chatPrompt: [{ role: "user", content: prompt }],
+      modelConfig: { temperature: 0.7 + 0.1 * candidateIndex }
+    });
+    if (response instanceof ReadableStream) {
+      return "";
+    }
+    return response.results[0]?.content || "";
+  }
+  /**
+   * Bootstraps few-shot examples for the program
+   */
+  async bootstrapFewShotExamples(metricFn) {
+    if (this.verbose) {
+      console.log("Bootstrapping few-shot examples...");
+    }
+    const result = await this.bootstrapper.compile(metricFn, {
+      maxDemos: this.maxBootstrappedDemos
+    });
+    return result.demos;
+  }
+  /**
+   * Selects labeled examples directly from the training set
+   */
+  selectLabeledExamples() {
+    const selectedExamples = [];
+    const indices = /* @__PURE__ */ new Set();
+    while (indices.size < this.maxLabeledDemos && indices.size < this.examples.length) {
+      const idx = Math.floor(Math.random() * this.examples.length);
+      if (!indices.has(idx)) {
+        indices.add(idx);
+        const example = this.examples[idx];
+        if (example) {
+          selectedExamples.push(example);
+        }
+      }
+    }
+    return selectedExamples;
+  }
+  /**
+   * Runs Bayesian optimization to find the best combination of few-shot examples and instructions
+   */
+  async runBayesianOptimization(bootstrappedDemos, labeledExamples, instructions, valset, metricFn) {
+    let bestConfig = null;
+    let bestScore = Number.NEGATIVE_INFINITY;
+    const evaluatedConfigs = [];
+    const defaultConfig = {
+      instruction: instructions[0] || "",
+      bootstrappedDemos: Math.min(1, bootstrappedDemos.length),
+      labeledExamples: Math.min(1, labeledExamples.length)
+    };
+    let trialsWithoutImprovement = 0;
+    let lastBestScore = Number.NEGATIVE_INFINITY;
+    const initialExplorationTrials = Math.min(
+      10,
+      Math.floor(this.numTrials / 3)
+    );
+    const configs = [];
+    for (let i = 0; i < initialExplorationTrials; i++) {
+      const instructionIndex = Math.floor(Math.random() * instructions.length);
+      const instructionValue = instructions[instructionIndex] || "";
+      const config = {
+        instruction: instructionValue,
+        bootstrappedDemos: Math.floor(
+          Math.random() * (bootstrappedDemos.length + 1)
+        ),
+        labeledExamples: Math.floor(
+          Math.random() * (labeledExamples.length + 1)
+        )
+      };
+      configs.push(config);
+    }
+    for (let i = 0; i < configs.length; i++) {
+      const config = configs[i];
+      if (!config) continue;
+      const score = await this.evaluateConfig(
+        config,
+        bootstrappedDemos,
+        labeledExamples,
+        valset,
+        metricFn,
+        i
+      );
+      evaluatedConfigs.push({ config, score });
+      if (score > bestScore) {
+        bestScore = score;
+        bestConfig = config;
+        if (this.verbose) {
+          console.log(
+            `New best configuration found with score ${bestScore} (exploration phase)`
+          );
+        }
+      }
+      updateProgressBar(
+        i + 1,
+        this.numTrials,
+        Math.round(bestScore * 100),
+        0,
+        "Running MIPROv2 optimization",
+        30
+      );
+    }
+    for (let i = configs.length; i < this.numTrials; i++) {
+      const nextConfig = this.selectNextConfiguration(
+        evaluatedConfigs,
+        bootstrappedDemos.length,
+        labeledExamples.length,
+        instructions
+      );
+      const score = await this.evaluateConfig(
+        nextConfig,
+        bootstrappedDemos,
+        labeledExamples,
+        valset,
+        metricFn,
+        i
+      );
+      evaluatedConfigs.push({ config: nextConfig, score });
+      if (score > bestScore) {
+        bestScore = score;
+        bestConfig = nextConfig;
+        if (this.verbose) {
+          console.log(
+            `New best configuration found with score ${bestScore} (exploitation phase)`
+          );
+        }
+        trialsWithoutImprovement = 0;
+        lastBestScore = bestScore;
+      } else {
+        if (bestScore - lastBestScore < this.minImprovementThreshold) {
+          trialsWithoutImprovement++;
+          if (trialsWithoutImprovement >= this.earlyStoppingTrials) {
+            if (this.verbose) {
+              console.log(
+                `Early stopping triggered after ${i + 1} trials. No improvement for ${trialsWithoutImprovement} trials.`
+              );
+            }
+            break;
+          }
+        } else {
+          lastBestScore = bestScore;
+          trialsWithoutImprovement = 0;
+        }
+      }
+      updateProgressBar(
+        i + 1,
+        this.numTrials,
+        Math.round(bestScore * 100),
+        0,
+        "Running MIPROv2 optimization",
+        30
+      );
+      if (this.minibatch && i > 0 && (i + 1) % this.minibatchFullEvalSteps === 0 && bestConfig) {
+        if (this.verbose) {
+          console.log(
+            `Running full evaluation on best configuration at trial ${i + 1}`
+          );
+        }
+        const fullScore = await this.fullEvaluation(
+          bestConfig,
+          bootstrappedDemos,
+          labeledExamples,
+          valset,
+          metricFn
+        );
+        if (this.verbose) {
+          console.log(`Full evaluation score: ${fullScore}`);
+        }
+        bestScore = fullScore;
+      }
+    }
+    if (!bestConfig) {
+      if (this.verbose) {
+        console.warn(
+          "Optimization failed to find any valid configurations, using default fallback configuration"
+        );
+      }
+      bestConfig = defaultConfig;
+      try {
+        bestScore = await this.evaluateConfig(
+          bestConfig,
+          bootstrappedDemos,
+          labeledExamples,
+          valset,
+          metricFn,
+          this.numTrials - 1
+        );
+      } catch (err) {
+        if (this.verbose) {
+          console.error("Error evaluating default configuration:", err);
+        }
+        bestScore = 0;
+      }
+    }
+    return { bestConfig, bestScore };
+  }
+  /**
+   * Evaluates a configuration on the validation set
+   */
+  async evaluateConfig(config, bootstrappedDemos, labeledExamples, valset, metricFn, trialIndex) {
+    this.applyConfigToProgram(
+      this.program,
+      config,
+      bootstrappedDemos,
+      labeledExamples
+    );
+    let evalSet = valset;
+    if (this.minibatch) {
+      const startIdx = trialIndex * this.minibatchSize % valset.length;
+      const minibatchEvalSet = [];
+      for (let j = 0; j < this.minibatchSize; j++) {
+        const idx = (startIdx + j) % valset.length;
+        const example = valset[idx];
+        if (example) {
+          minibatchEvalSet.push(example);
+        }
+      }
+      evalSet = minibatchEvalSet;
+    }
+    let correctCount = 0;
+    for (const example of evalSet) {
+      try {
+        const prediction = await this.program.forward(this.ai, example);
+        const correct = metricFn({ prediction, example });
+        if (correct) correctCount++;
+      } catch (err) {
+        if (this.verbose) {
+          console.error("Error evaluating example:", err);
+        }
+      }
+    }
+    return correctCount / evalSet.length;
+  }
+  /**
+   * Run full evaluation on the entire validation set
+   */
+  async fullEvaluation(config, bootstrappedDemos, labeledExamples, valset, metricFn) {
+    this.applyConfigToProgram(
+      this.program,
+      config,
+      bootstrappedDemos,
+      labeledExamples
+    );
+    let fullCorrectCount = 0;
+    for (const example of valset) {
+      try {
+        const prediction = await this.program.forward(this.ai, example);
+        const correct = metricFn({ prediction, example });
+        if (correct) fullCorrectCount++;
+      } catch (err) {
+        if (this.verbose) {
+          console.error("Error evaluating example:", err);
+        }
+      }
+    }
+    return fullCorrectCount / valset.length;
+  }
+  /**
+   * Implements a Bayesian-inspired selection of the next configuration to try
+   * This is a simplified version using Upper Confidence Bound (UCB) strategy
+   */
+  selectNextConfiguration(evaluatedConfigs, maxBootstrappedDemos, maxLabeledExamples, instructions) {
+    if (evaluatedConfigs.length < 5) {
+      const instructionIndex = Math.floor(Math.random() * instructions.length);
+      return {
+        instruction: instructions[instructionIndex] || "",
+        bootstrappedDemos: Math.floor(
+          Math.random() * (maxBootstrappedDemos + 1)
+        ),
+        labeledExamples: Math.floor(Math.random() * (maxLabeledExamples + 1))
+      };
+    }
+    const sortedConfigs = [...evaluatedConfigs].sort(
+      (a, b) => b.score - a.score
+    );
+    const topConfigs = sortedConfigs.slice(0, Math.min(3, sortedConfigs.length));
+    const meanBootstrappedDemos = topConfigs.reduce((sum, c) => sum + c.config.bootstrappedDemos, 0) / topConfigs.length;
+    const meanLabeledExamples = topConfigs.reduce((sum, c) => sum + c.config.labeledExamples, 0) / topConfigs.length;
+    const popularInstructions = topConfigs.map((c) => c.config.instruction);
+    const explorationFactor = Math.max(
+      0.2,
+      1 - evaluatedConfigs.length / this.numTrials
+    );
+    let newBootstrappedDemos;
+    let newLabeledExamples;
+    let newInstruction;
+    if (Math.random() < 0.7) {
+      newBootstrappedDemos = Math.min(
+        maxBootstrappedDemos,
+        Math.max(
+          0,
+          Math.round(
+            meanBootstrappedDemos + (Math.random() * 2 - 1) * explorationFactor * 2
+          )
+        )
+      );
+    } else {
+      newBootstrappedDemos = Math.floor(
+        Math.random() * (maxBootstrappedDemos + 1)
+      );
+    }
+    if (Math.random() < 0.7) {
+      newLabeledExamples = Math.min(
+        maxLabeledExamples,
+        Math.max(
+          0,
+          Math.round(
+            meanLabeledExamples + (Math.random() * 2 - 1) * explorationFactor * 2
+          )
+        )
+      );
+    } else {
+      newLabeledExamples = Math.floor(Math.random() * (maxLabeledExamples + 1));
+    }
+    if (Math.random() < 0.7 && popularInstructions.length > 0) {
+      const idx = Math.floor(Math.random() * popularInstructions.length);
+      newInstruction = popularInstructions[idx] || "";
+    } else {
+      const idx = Math.floor(Math.random() * instructions.length);
+      newInstruction = instructions[idx] || "";
+    }
+    return {
+      instruction: newInstruction,
+      bootstrappedDemos: newBootstrappedDemos,
+      labeledExamples: newLabeledExamples
+    };
+  }
+  /**
+   * Applies a configuration to a program instance
+   */
+  applyConfigToProgram(program, config, bootstrappedDemos, labeledExamples) {
+    this.setInstructionToProgram(program, config.instruction);
+    if (config.bootstrappedDemos > 0) {
+      program.setDemos(bootstrappedDemos.slice(0, config.bootstrappedDemos));
+    }
+    if (config.labeledExamples > 0) {
+      program.setExamples(labeledExamples.slice(0, config.labeledExamples));
+    }
+  }
+  /**
+   * Sets instruction to a program
+   * Note: Workaround since setInstruction may not be available directly
+   */
+  setInstructionToProgram(program, instruction) {
+    const programWithInstruction = program;
+    programWithInstruction.setInstruction?.(instruction);
+  }
+  /**
+   * The main compile method to run MIPROv2 optimization
+   * @param metricFn Evaluation metric function
+   * @param options Optional configuration options
+   * @returns The optimized program
+   */
+  async compile(metricFn, options) {
+    if (options?.auto) {
+      this.configureAuto(options.auto);
+    }
+    const trainset = this.examples;
+    const valset = options?.valset || this.examples.slice(0, Math.floor(this.examples.length * 0.8));
+    if (this.verbose) {
+      console.log(`Starting MIPROv2 optimization with ${this.numTrials} trials`);
+      console.log(
+        `Using ${trainset.length} examples for training and ${valset.length} for validation`
+      );
+    }
+    if (options?.teacher) {
+      if (this.verbose) {
+        console.log("Using provided teacher to assist with bootstrapping");
+      }
+      const bootstrapperWithTeacher = new AxBootstrapFewShot({
+        ai: this.ai,
+        program: this.program,
+        examples: this.examples,
+        options: {
+          maxDemos: this.maxBootstrappedDemos,
+          maxRounds: 3,
+          verboseMode: this.verbose,
+          teacherAI: this.ai
+          // Use the same AI but with the teacher program
+        }
+      });
+      this.bootstrapper = bootstrapperWithTeacher;
+    }
+    let bootstrappedDemos = [];
+    if (this.maxBootstrappedDemos > 0) {
+      bootstrappedDemos = await this.bootstrapFewShotExamples(metricFn);
+      if (this.verbose) {
+        console.log(
+          `Generated ${bootstrappedDemos.length} bootstrapped demonstrations`
+        );
+      }
+    }
+    let labeledExamples = [];
+    if (this.maxLabeledDemos > 0) {
+      labeledExamples = this.selectLabeledExamples();
+      if (this.verbose) {
+        console.log(
+          `Selected ${labeledExamples.length} labeled examples from training set`
+        );
+      }
+    }
+    const instructions = await this.proposeInstructionCandidates();
+    if (this.verbose) {
+      console.log(`Generated ${instructions.length} instruction candidates`);
+    }
+    const { bestConfig, bestScore } = await this.runBayesianOptimization(
+      bootstrappedDemos,
+      labeledExamples,
+      instructions,
+      valset,
+      metricFn
+    );
+    if (this.verbose) {
+      console.log(`Optimization complete. Best score: ${bestScore}`);
+      console.log(`Best configuration: ${JSON.stringify(bestConfig)}`);
+    }
+    this.applyConfigToProgram(
+      this.program,
+      bestConfig,
+      bootstrappedDemos,
+      labeledExamples
+    );
+    return this.program;
+  }
+};
+
 // ai/mock/api.ts
 var AxMockAIService = class {
   constructor(config = {}) {
@@ -10378,6 +11179,7 @@ export {
   AxMCPHTTPTransport,
   AxMCPStdioTransport,
   AxMemory,
+  AxMiPRO,
   AxMockAIService,
   AxMultiServiceRouter,
   AxProgram,