@ax-llm/ax 12.0.7 → 12.0.8

package/index.cjs

@@ -81,6 +81,7 @@ __export(index_exports, {
   AxAssertionError: () => AxAssertionError,
   AxBalancer: () => AxBalancer,
   AxBaseAI: () => AxBaseAI,
+  AxBaseOptimizer: () => AxBaseOptimizer,
   AxBootstrapFewShot: () => AxBootstrapFewShot,
   AxChainOfThought: () => AxChainOfThought,
   AxDB: () => AxDB,
@@ -90,6 +91,7 @@ __export(index_exports, {
   AxDBMemory: () => AxDBMemory,
   AxDBPinecone: () => AxDBPinecone,
   AxDBWeaviate: () => AxDBWeaviate,
+  AxDefaultCostTracker: () => AxDefaultCostTracker,
   AxDefaultQueryRewriter: () => AxDefaultQueryRewriter,
   AxDefaultResultReranker: () => AxDefaultResultReranker,
   AxDockerSession: () => AxDockerSession,
@@ -7538,8 +7540,9 @@ var AxInstanceRegistry = class {
     this.reg.add(instance);
   }
   *[Symbol.iterator]() {
-    for (const key of this.reg) {
-      yield key;
+    const items = Array.from(this.reg);
+    for (let i = 0; i < items.length; i++) {
+      yield items[i];
     }
   }
 };
@@ -8477,7 +8480,7 @@ var AxSignature = class _AxSignature {
       this.getOutputFields().forEach((field) => {
         validateField(field, "output");
       });
-      this.sigHash = (0, import_crypto3.createHash)("sha256").update(this.description ?? "").update(JSON.stringify(this.inputFields)).update(JSON.stringify(this.outputFields)).digest("hex");
+      this.sigHash = (0, import_crypto3.createHash)("sha256").update(JSON.stringify(this.inputFields)).update(JSON.stringify(this.outputFields)).digest("hex");
       this.sigString = renderSignature(
         this.description,
         this.inputFields,
@@ -8798,7 +8801,7 @@ var AxProgramWithSignature = class {
     this.signature.validate();
     this.sigHash = this.signature?.hash();
     this.children = new AxInstanceRegistry();
-    this.key = { id: this.constructor.name };
+    this.key = { id: this.signature.hash() };
   }
   getSignature() {
     return this.signature;
@@ -8818,8 +8821,8 @@ var AxProgramWithSignature = class {
   }
   setId(id) {
     this.key = { id, custom: true };
-    for (const child of this.children) {
-      child.setParentId(id);
+    for (const child of Array.from(this.children)) {
+      child?.setParentId(id);
     }
   }
   setParentId(parentId) {
@@ -8832,8 +8835,8 @@ var AxProgramWithSignature = class {
     if (!("programId" in examples)) {
       return;
     }
-    for (const child of this.children) {
-      child.setExamples(examples, options);
+    for (const child of Array.from(this.children)) {
+      child?.setExamples(examples, options);
     }
   }
   _setExamples(examples, options) {
@@ -8866,30 +8869,37 @@ var AxProgramWithSignature = class {
     if (this.trace) {
       traces.push({ trace: this.trace, programId: this.key.id });
     }
-    for (const child of this.children) {
-      const _traces = child.getTraces();
-      traces = [...traces, ..._traces];
+    for (const child of Array.from(this.children)) {
+      const _traces = child?.getTraces();
+      traces = [...traces, ..._traces ?? []];
     }
     return traces;
   }
   getUsage() {
     let usage = [...this.usage ?? []];
-    for (const child of this.children) {
-      const cu = child.getUsage();
-      usage = [...usage, ...cu];
+    for (const child of Array.from(this.children)) {
+      const cu = child?.getUsage();
+      usage = [...usage, ...cu ?? []];
     }
     return mergeProgramUsage(usage);
   }
   resetUsage() {
     this.usage = [];
-    for (const child of this.children) {
-      child.resetUsage();
+    for (const child of Array.from(this.children)) {
+      child?.resetUsage();
     }
   }
   setDemos(demos) {
+    const hasChildren = Array.from(this.children).length > 0;
+    const hasMatchingDemo = demos.some((demo) => demo.programId === this.key.id);
+    if (hasChildren && !hasMatchingDemo) {
+      throw new Error(
+        `Program with id '${this.key.id}' has children but no matching programId found in demos`
+      );
+    }
     this.demos = demos.filter((v) => v.programId === this.key.id).map((v) => v.traces).flat();
-    for (const child of this.children) {
-      child.setDemos(demos);
+    for (const child of Array.from(this.children)) {
+      child?.setDemos(demos);
     }
   }
 };
@@ -8917,8 +8927,8 @@ var AxProgram = class {
   }
   setId(id) {
     this.key = { id, custom: true };
-    for (const child of this.children) {
-      child.setParentId(id);
+    for (const child of Array.from(this.children)) {
+      child?.setParentId(id);
     }
   }
   setParentId(parentId) {
@@ -8930,8 +8940,8 @@ var AxProgram = class {
     if (!("programId" in examples)) {
       return;
     }
-    for (const child of this.children) {
-      child.setExamples(examples, options);
+    for (const child of Array.from(this.children)) {
+      child?.setExamples(examples, options);
     }
   }
   getTraces() {
@@ -8939,29 +8949,36 @@ var AxProgram = class {
     if (this.trace) {
       traces.push({ trace: this.trace, programId: this.key.id });
     }
-    for (const child of this.children) {
-      const _traces = child.getTraces();
-      traces = [...traces, ..._traces];
+    for (const child of Array.from(this.children)) {
+      const _traces = child?.getTraces();
+      traces = [...traces, ..._traces ?? []];
     }
     return traces;
   }
   getUsage() {
     let usage = [...this.usage ?? []];
-    for (const child of this.children) {
-      const cu = child.getUsage();
-      usage = [...usage, ...cu];
+    for (const child of Array.from(this.children)) {
+      const cu = child?.getUsage();
+      usage = [...usage, ...cu ?? []];
     }
     return mergeProgramUsage(usage);
   }
   resetUsage() {
     this.usage = [];
-    for (const child of this.children) {
-      child.resetUsage();
+    for (const child of Array.from(this.children)) {
+      child?.resetUsage();
     }
   }
   setDemos(demos) {
-    for (const child of this.children) {
-      child.setDemos(demos);
+    const hasChildren = Array.from(this.children).length > 0;
+    const hasMatchingDemo = demos.some((demo) => demo.programId === this.key.id);
+    if (hasChildren && !hasMatchingDemo) {
+      throw new Error(
+        `Program with id '${this.key.id}' has children but no matching programId found in demos`
+      );
+    }
+    for (const child of Array.from(this.children)) {
+      child?.setDemos(demos);
     }
   }
 };
@@ -9689,7 +9706,9 @@ var AxAgent = class {
       description: definition ?? description
     });
     for (const agent of agents ?? []) {
-      this.program.register(agent);
+      this.program.register(
+        agent
+      );
     }
     this.name = name;
     this.func = {
@@ -10193,6 +10212,673 @@ function validateModels2(services) {
   }
 }
 
+// dsp/optimizer.ts
+var AxDefaultCostTracker = class {
+  tokenUsage = {};
+  totalTokens = 0;
+  // Configuration options
+  costPerModel;
+  maxCost;
+  maxTokens;
+  constructor(options) {
+    this.costPerModel = options?.costPerModel ?? {};
+    this.maxCost = options?.maxCost;
+    this.maxTokens = options?.maxTokens;
+  }
+  trackTokens(count, model) {
+    this.tokenUsage[model] = (this.tokenUsage[model] || 0) + count;
+    this.totalTokens += count;
+  }
+  getCurrentCost() {
+    let totalCost = 0;
+    for (const [model, tokens] of Object.entries(this.tokenUsage)) {
+      const costPer1K = this.costPerModel[model] || 1e-3;
+      totalCost += tokens / 1e3 * costPer1K;
+    }
+    return totalCost;
+  }
+  getTokenUsage() {
+    return { ...this.tokenUsage };
+  }
+  getTotalTokens() {
+    return this.totalTokens;
+  }
+  isLimitReached() {
+    if (this.maxTokens !== void 0 && this.totalTokens >= this.maxTokens) {
+      return true;
+    }
+    if (this.maxCost !== void 0) {
+      const currentCost = this.getCurrentCost();
+      if (currentCost >= this.maxCost) {
+        return true;
+      }
+    }
+    return false;
+  }
+  reset() {
+    this.tokenUsage = {};
+    this.totalTokens = 0;
+  }
+};
+var AxBaseOptimizer = class {
+  // Common AxOptimizerArgs fields
+  studentAI;
+  teacherAI;
+  examples;
+  validationSet;
+  targetScore;
+  minSuccessRate;
+  onProgress;
+  onEarlyStop;
+  costTracker;
+  seed;
+  // Checkpointing fields
+  checkpointSave;
+  checkpointLoad;
+  checkpointInterval;
+  resumeFromCheckpoint;
+  // Checkpoint state
+  currentRound = 0;
+  scoreHistory = [];
+  configurationHistory = [];
+  // Common optimization statistics
+  stats;
+  constructor(args) {
+    if (args.examples.length === 0) {
+      throw new Error("No examples found");
+    }
+    this.studentAI = args.studentAI;
+    this.teacherAI = args.teacherAI;
+    this.examples = args.examples;
+    this.validationSet = args.validationSet;
+    this.targetScore = args.targetScore;
+    this.minSuccessRate = args.minSuccessRate;
+    this.onProgress = args.onProgress;
+    this.onEarlyStop = args.onEarlyStop;
+    this.seed = args.seed;
+    this.checkpointSave = args.checkpointSave;
+    this.checkpointLoad = args.checkpointLoad;
+    this.checkpointInterval = args.checkpointInterval ?? 10;
+    this.resumeFromCheckpoint = args.resumeFromCheckpoint;
+    const costTracker = new AxDefaultCostTracker({
+      maxTokens: 1e6
+    });
+    this.costTracker = args.costTracker ?? costTracker;
+    this.stats = this.initializeStats();
+  }
+  /**
+   * Initialize the optimization statistics structure
+   */
+  initializeStats() {
+    return {
+      totalCalls: 0,
+      successfulDemos: 0,
+      estimatedTokenUsage: 0,
+      earlyStopped: false,
+      resourceUsage: {
+        totalTokens: 0,
+        totalTime: 0,
+        avgLatencyPerEval: 0,
+        costByModel: {}
+      },
+      convergenceInfo: {
+        converged: false,
+        finalImprovement: 0,
+        stagnationRounds: 0,
+        convergenceThreshold: 0.01
+      }
+    };
+  }
+  /**
+   * Set up reproducible random seed if provided
+   */
+  setupRandomSeed() {
+    if (this.seed !== void 0) {
+      Math.random = (() => {
+        let seed = this.seed;
+        return () => {
+          seed = (seed * 9301 + 49297) % 233280;
+          return seed / 233280;
+        };
+      })();
+    }
+  }
+  /**
+   * Check if optimization should stop early due to cost limits
+   */
+  checkCostLimits() {
+    return this.costTracker?.isLimitReached() ?? false;
+  }
+  /**
+   * Check if target score has been reached
+   */
+  checkTargetScore(currentScore) {
+    return this.targetScore !== void 0 && currentScore >= this.targetScore;
+  }
+  /**
+   * Update resource usage statistics
+   */
+  updateResourceUsage(startTime, tokensUsed = 0) {
+    this.stats.resourceUsage.totalTime = Date.now() - startTime;
+    this.stats.resourceUsage.totalTokens += tokensUsed;
+    if (this.stats.totalCalls > 0) {
+      this.stats.resourceUsage.avgLatencyPerEval = this.stats.resourceUsage.totalTime / this.stats.totalCalls;
+    }
+  }
+  /**
+   * Trigger early stopping with appropriate callbacks
+   */
+  triggerEarlyStopping(reason, bestScoreRound) {
+    this.stats.earlyStopped = true;
+    this.stats.earlyStopping = {
+      bestScoreRound,
+      patienceExhausted: reason.includes("improvement"),
+      reason
+    };
+    if (this.onEarlyStop) {
+      this.onEarlyStop(reason, this.stats);
+    }
+  }
+  /**
+   * Get the validation set, with fallback to a split of examples
+   */
+  getValidationSet(options) {
+    return options?.overrideValidationSet || this.validationSet || this.examples.slice(0, Math.floor(this.examples.length * 0.2));
+  }
+  /**
+   * Get the AI service to use for a specific task, preferring teacher when available
+   * @param preferTeacher Whether to prefer teacher AI over student AI
+   * @param options Optional compile options that may override teacher AI
+   * @returns The appropriate AI service to use
+   */
+  getAIService(preferTeacher = false, options) {
+    if (preferTeacher && options?.overrideTeacherAI) {
+      return options.overrideTeacherAI;
+    }
+    if (preferTeacher && this.teacherAI) {
+      return this.teacherAI;
+    }
+    return this.studentAI;
+  }
+  /**
+   * Check if teacher AI is available (including overrides)
+   * @param options Optional compile options that may override teacher AI
+   * @returns True if teacher AI is configured or overridden
+   */
+  hasTeacherAI(options) {
+    return options?.overrideTeacherAI !== void 0 || this.teacherAI !== void 0;
+  }
+  /**
+   * Get teacher AI if available, otherwise return student AI
+   * @param options Optional compile options that may override teacher AI
+   * @returns Teacher AI if available, otherwise student AI
+   */
+  getTeacherOrStudentAI(options) {
+    return options?.overrideTeacherAI || this.teacherAI || this.studentAI;
+  }
+  /**
+   * Execute a task with teacher AI if available, otherwise use student AI
+   * @param task Function that takes an AI service and returns a promise
+   * @param preferTeacher Whether to prefer teacher AI (default: true)
+   * @param options Optional compile options that may override teacher AI
+   * @returns Result of the task execution
+   */
+  async executeWithTeacher(task, preferTeacher = true, options) {
+    const ai = this.getAIService(preferTeacher, options);
+    return await task(ai);
+  }
+  /**
+   * Get current optimization statistics
+   */
+  getStats() {
+    return { ...this.stats };
+  }
+  /**
+   * Reset optimizer state for reuse with different programs
+   */
+  reset() {
+    this.stats = this.initializeStats();
+    this.costTracker?.reset();
+    this.currentRound = 0;
+    this.scoreHistory = [];
+    this.configurationHistory = [];
+  }
+  /**
+   * Basic program validation that can be extended by concrete optimizers
+   */
+  validateProgram(program) {
+    const issues = [];
+    const suggestions = [];
+    if (!("forward" in program) || typeof program.forward !== "function") {
+      issues.push("Program must have a forward method");
+    }
+    if (this.examples.length < 2) {
+      issues.push("Need at least 2 examples for optimization");
+      suggestions.push("Provide more training examples");
+    }
+    const valSetSize = this.getValidationSet().length;
+    if (valSetSize < 1) {
+      issues.push("Validation set is empty");
+      suggestions.push("Provide examples or a validation set");
+    }
+    return {
+      isValid: issues.length === 0,
+      issues,
+      suggestions
+    };
+  }
+  /**
+   * Multi-objective optimization using Pareto frontier
+   * Default implementation that leverages the single-objective compile method
+   * @param program The program to optimize
+   * @param metricFn Multi-objective metric function that returns multiple scores
+   * @param options Optional configuration options
+   * @returns Pareto optimization result with frontier of non-dominated solutions
+   */
+  async compilePareto(program, metricFn, options) {
+    const startTime = Date.now();
+    if (options?.verbose) {
+      console.log("Starting Pareto optimization using base implementation");
+      console.log("This will run multiple single-objective optimizations");
+    }
+    const solutions = await this.generateWeightedSolutions(
+      program,
+      metricFn,
+      options
+    );
+    const constraintSolutions = await this.generateConstraintSolutions(
+      program,
+      metricFn,
+      options
+    );
+    const allSolutions = [...solutions, ...constraintSolutions];
+    if (options?.verbose) {
+      console.log(`Generated ${allSolutions.length} candidate solutions`);
+    }
+    const paretoFront = this.findParetoFrontier(allSolutions);
+    const hypervolume = this.calculateHypervolume(paretoFront);
+    if (options?.verbose) {
+      console.log(`Found ${paretoFront.length} non-dominated solutions`);
+      console.log(`Hypervolume: ${hypervolume?.toFixed(4) || "N/A"}`);
+    }
+    this.updateResourceUsage(startTime);
+    this.stats.convergenceInfo.converged = true;
+    const bestScore = paretoFront.length > 0 ? Math.max(
+      ...paretoFront.map((sol) => Math.max(...Object.values(sol.scores)))
+    ) : 0;
+    return {
+      demos: paretoFront.length > 0 ? [...paretoFront[0].demos] : void 0,
+      stats: this.stats,
+      bestScore,
+      paretoFront,
+      hypervolume,
+      paretoFrontSize: paretoFront.length,
+      finalConfiguration: {
+        paretoFrontSize: paretoFront.length,
+        hypervolume,
+        strategy: "weighted_combinations_and_constraints",
+        numSolutions: allSolutions.length
+      }
+    };
+  }
+  /**
+   * Generate solutions using different weighted combinations of objectives
+   */
+  async generateWeightedSolutions(program, metricFn, options) {
+    const solutions = [];
+    const sampleExample = this.examples[0];
+    const samplePrediction = await program.forward(
+      this.studentAI,
+      sampleExample
+    );
+    const sampleScores = await metricFn({
+      prediction: samplePrediction,
+      example: sampleExample
+    });
+    const objectives = Object.keys(sampleScores);
+    if (options?.verbose) {
+      console.log(`Detected objectives: ${objectives.join(", ")}`);
+    }
+    const weightCombinations = this.generateWeightCombinations(objectives);
+    for (let i = 0; i < weightCombinations.length; i++) {
+      const weights = weightCombinations[i];
+      if (options?.verbose) {
+        console.log(`Optimizing with weights: ${JSON.stringify(weights)}`);
+      }
+      const weightedMetric = async ({ prediction, example }) => {
+        const scores = await metricFn({ prediction, example });
+        let weightedScore = 0;
+        for (const [objective, score] of Object.entries(scores)) {
+          weightedScore += score * (weights[objective] || 0);
+        }
+        return weightedScore;
+      };
+      try {
+        const result = await this.compile(program, weightedMetric, {
+          ...options,
+          verbose: false
+          // Suppress inner optimization logs
+        });
+        const scores = await this.evaluateWithMultiObjective(
+          program,
+          result,
+          metricFn
+        );
+        solutions.push({
+          scores,
+          demos: result.demos,
+          configuration: {
+            ...result.finalConfiguration,
+            weights,
+            strategy: "weighted_combination"
+          }
+        });
+      } catch (error) {
+        if (options?.verbose) {
+          console.warn(
+            `Failed optimization with weights ${JSON.stringify(weights)}:`,
+            error
+          );
+        }
+        continue;
+      }
+    }
+    return solutions;
+  }
+  /**
+   * Generate solutions using constraint-based optimization
+   */
+  async generateConstraintSolutions(program, metricFn, options) {
+    const solutions = [];
+    const sampleExample = this.examples[0];
+    const samplePrediction = await program.forward(
+      this.studentAI,
+      sampleExample
+    );
+    const sampleScores = await metricFn({
+      prediction: samplePrediction,
+      example: sampleExample
+    });
+    const objectives = Object.keys(sampleScores);
+    for (const primaryObjective of objectives) {
+      if (options?.verbose) {
+        console.log(
+          `Optimizing ${primaryObjective} with constraints on other objectives`
+        );
+      }
+      const constraintMetric = async ({ prediction, example }) => {
+        const scores = await metricFn({ prediction, example });
+        const primaryScore = scores[primaryObjective] || 0;
+        let penalty = 0;
+        for (const [objective, score] of Object.entries(scores)) {
+          if (objective !== primaryObjective) {
+            if (score < 0.3) {
+              penalty += (0.3 - score) * 2;
+            }
+          }
+        }
+        return primaryScore - penalty;
+      };
+      try {
+        const result = await this.compile(program, constraintMetric, {
+          ...options,
+          verbose: false
+        });
+        const scores = await this.evaluateWithMultiObjective(
+          program,
+          result,
+          metricFn
+        );
+        solutions.push({
+          scores,
+          demos: result.demos,
+          configuration: {
+            ...result.finalConfiguration,
+            primaryObjective,
+            strategy: "constraint_based"
+          }
+        });
+      } catch (error) {
+        if (options?.verbose) {
+          console.warn(
+            `Failed constraint optimization for ${primaryObjective}:`,
+            error
+          );
+        }
+        continue;
+      }
+    }
+    return solutions;
+  }
+  /**
+   * Generate different weight combinations for objectives
+   */
+  generateWeightCombinations(objectives) {
+    const combinations = [];
+    for (const objective of objectives) {
+      const weights = {};
+      for (const obj of objectives) {
+        weights[obj] = obj === objective ? 1 : 0;
+      }
+      combinations.push(weights);
+    }
+    const equalWeights = {};
+    for (const objective of objectives) {
+      equalWeights[objective] = 1 / objectives.length;
+    }
+    combinations.push(equalWeights);
+    if (objectives.length === 2) {
+      const [obj1, obj2] = objectives;
+      for (let w1 = 0.1; w1 <= 0.9; w1 += 0.2) {
+        const w2 = 1 - w1;
+        combinations.push({ [obj1]: w1, [obj2]: w2 });
+      }
+    }
+    if (objectives.length === 3) {
+      const [obj1, obj2, obj3] = objectives;
+      combinations.push(
+        { [obj1]: 0.5, [obj2]: 0.3, [obj3]: 0.2 },
+        { [obj1]: 0.3, [obj2]: 0.5, [obj3]: 0.2 },
+        { [obj1]: 0.2, [obj2]: 0.3, [obj3]: 0.5 }
+      );
+    }
+    return combinations;
+  }
+  /**
+   * Evaluate a single-objective result with multi-objective metrics
+   */
+  async evaluateWithMultiObjective(program, result, metricFn) {
+    const valSet = this.getValidationSet();
+    const allScores = {};
+    const testProgram = { ...program };
+    if (result.demos && "setDemos" in testProgram) {
+      ;
+      testProgram.setDemos(result.demos);
+    }
+    const evalSet = valSet.slice(0, Math.min(5, valSet.length));
+    for (const example of evalSet) {
+      try {
+        const prediction = await testProgram.forward(
+          this.studentAI,
+          example
+        );
+        const scores = await metricFn({ prediction, example });
+        for (const [objective, score] of Object.entries(scores)) {
+          if (!allScores[objective]) {
+            allScores[objective] = [];
+          }
+          allScores[objective].push(score);
+        }
+      } catch {
+        continue;
+      }
+    }
+    const avgScores = {};
+    for (const [objective, scores] of Object.entries(allScores)) {
+      avgScores[objective] = scores.length > 0 ? scores.reduce((sum, score) => sum + score, 0) / scores.length : 0;
+    }
+    return avgScores;
+  }
+  /**
+   * Find the Pareto frontier from a set of solutions
+   */
+  findParetoFrontier(solutions) {
+    const paretoFront = [];
+    for (let i = 0; i < solutions.length; i++) {
+      const solutionA = solutions[i];
+      let isDominated = false;
+      let dominatedCount = 0;
+      for (let j = 0; j < solutions.length; j++) {
+        if (i === j) continue;
+        const solutionB = solutions[j];
+        if (this.dominates(solutionB.scores, solutionA.scores)) {
+          isDominated = true;
+          break;
+        }
+        if (this.dominates(solutionA.scores, solutionB.scores)) {
+          dominatedCount++;
+        }
+      }
+      if (!isDominated) {
+        paretoFront.push({
+          demos: solutionA.demos || [],
+          scores: solutionA.scores,
+          configuration: solutionA.configuration,
+          dominatedSolutions: dominatedCount
+        });
+      }
+    }
+    return paretoFront;
+  }
+  /**
+   * Check if solution A dominates solution B
+   * A dominates B if A is better or equal in all objectives and strictly better in at least one
+   */
+  dominates(scoresA, scoresB) {
+    const objectives = Object.keys(scoresA);
+    let atLeastAsGood = true;
+    let strictlyBetter = false;
+    for (const objective of objectives) {
+      const scoreA = scoresA[objective] || 0;
+      const scoreB = scoresB[objective] || 0;
+      if (scoreA < scoreB) {
+        atLeastAsGood = false;
+        break;
+      }
+      if (scoreA > scoreB) {
+        strictlyBetter = true;
+      }
+    }
+    return atLeastAsGood && strictlyBetter;
+  }
+  /**
+   * Calculate hypervolume of the Pareto frontier
+   * Simplified implementation using reference point at origin
+   */
+  calculateHypervolume(paretoFront) {
+    if (paretoFront.length === 0) return void 0;
+    const firstSolution = paretoFront[0];
+    const objectives = Object.keys(firstSolution.scores);
+    if (objectives.length === 2) {
+      const [obj1, obj2] = objectives;
+      let hypervolume = 0;
+      const sortedSolutions = [...paretoFront].sort(
+        (a, b) => (b.scores[obj1] || 0) - (a.scores[obj1] || 0)
+      );
+      let prevScore2 = 0;
+      for (const solution of sortedSolutions) {
+        const score1 = solution.scores[obj1] || 0;
+        const score2 = solution.scores[obj2] || 0;
+        hypervolume += score1 * (score2 - prevScore2);
+        prevScore2 = Math.max(prevScore2, score2);
+      }
+      return hypervolume;
+    }
+    return void 0;
+  }
+  /**
+   * Save current optimization state to checkpoint
+   */
+  async saveCheckpoint(optimizerType, optimizerConfig, bestScore, bestConfiguration, optimizerState = {}, options) {
+    const saveFn = options?.overrideCheckpointSave || this.checkpointSave;
+    if (!saveFn) return void 0;
+    const checkpoint = {
+      version: "1.0.0",
+      timestamp: Date.now(),
+      optimizerType,
+      optimizerConfig,
+      currentRound: this.currentRound,
+      totalRounds: this.stats.resourceUsage.totalTime > 0 ? this.currentRound : 0,
+      bestScore,
+      bestConfiguration,
+      scoreHistory: [...this.scoreHistory],
+      configurationHistory: [...this.configurationHistory],
+      stats: { ...this.stats },
+      optimizerState,
+      examples: this.examples,
+      validationSet: this.validationSet
+    };
+    return await saveFn(checkpoint);
+  }
+  /**
+   * Load optimization state from checkpoint
+   */
+  async loadCheckpoint(checkpointId, options) {
+    const loadFn = options?.overrideCheckpointLoad || this.checkpointLoad;
+    if (!loadFn) return null;
+    return await loadFn(checkpointId);
+  }
+  /**
+   * Restore optimizer state from checkpoint
+   */
+  restoreFromCheckpoint(checkpoint) {
+    this.currentRound = checkpoint.currentRound;
+    this.scoreHistory = [...checkpoint.scoreHistory];
+    this.configurationHistory = [...checkpoint.configurationHistory];
+    this.stats = { ...checkpoint.stats };
+  }
+  /**
+   * Check if checkpoint should be saved
+   */
+  shouldSaveCheckpoint(round, options) {
+    const interval = options?.overrideCheckpointInterval || this.checkpointInterval;
+    return interval !== void 0 && round % interval === 0;
+  }
+  /**
+   * Update optimization progress and handle checkpointing
+   */
+  async updateOptimizationProgress(round, score, configuration, optimizerType, optimizerConfig, bestScore, bestConfiguration, optimizerState = {}, options) {
+    this.currentRound = round;
+    this.scoreHistory.push(score);
+    this.configurationHistory.push(configuration);
+    if (this.shouldSaveCheckpoint(round, options)) {
+      await this.saveCheckpoint(
+        optimizerType,
+        optimizerConfig,
+        bestScore,
+        bestConfiguration,
+        optimizerState,
+        options
+      );
+    }
+  }
+  /**
+   * Save final checkpoint on completion
+   */
+  async saveFinalCheckpoint(optimizerType, optimizerConfig, bestScore, bestConfiguration, optimizerState = {}, options) {
+    if (options?.saveCheckpointOnComplete !== false) {
+      await this.saveCheckpoint(
+        optimizerType,
+        optimizerConfig,
+        bestScore,
+        bestConfiguration,
+        { ...optimizerState, final: true },
+        options
+      );
+    }
+  }
+};
+
 // db/base.ts
 var import_api23 = require("@opentelemetry/api");
 var AxDBBase = class {
@@ -11652,11 +12338,7 @@ var AxMCPStreambleHTTPTransport = class {
 };
 
 // dsp/optimizers/bootstrapFewshot.ts
-var AxBootstrapFewShot = class {
-  ai;
-  teacherAI;
-  program;
-  examples;
+var AxBootstrapFewShot = class extends AxBaseOptimizer {
   maxRounds;
   maxDemos;
   maxExamples;
@@ -11667,37 +12349,20 @@ var AxBootstrapFewShot = class {
   verboseMode;
   debugMode;
   traces = [];
-  stats = {
-    totalCalls: 0,
-    successfulDemos: 0,
-    estimatedTokenUsage: 0,
-    earlyStopped: false
-  };
-  constructor({
-    ai,
-    program,
-    examples = [],
-    options
-  }) {
-    if (examples.length === 0) {
-      throw new Error("No examples found");
-    }
-    const bootstrapOptions = options;
-    this.maxRounds = bootstrapOptions?.maxRounds ?? 3;
-    this.maxDemos = bootstrapOptions?.maxDemos ?? 4;
-    this.maxExamples = bootstrapOptions?.maxExamples ?? 16;
-    this.batchSize = bootstrapOptions?.batchSize ?? 1;
-    this.earlyStoppingPatience = bootstrapOptions?.earlyStoppingPatience ?? 0;
-    this.costMonitoring = bootstrapOptions?.costMonitoring ?? false;
-    this.maxTokensPerGeneration = bootstrapOptions?.maxTokensPerGeneration ?? 0;
-    this.verboseMode = bootstrapOptions?.verboseMode ?? true;
-    this.debugMode = bootstrapOptions?.debugMode ?? false;
-    this.ai = ai;
-    this.teacherAI = bootstrapOptions?.teacherAI;
-    this.program = program;
-    this.examples = examples;
-  }
-  async compileRound(roundIndex, metricFn, options) {
+  constructor(args) {
+    super(args);
+    const options = args.options || {};
+    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;
+  }
+  async compileRound(program, roundIndex, metricFn, options) {
     const st = (/* @__PURE__ */ new Date()).getTime();
     const maxDemos = options?.maxDemos ?? this.maxDemos;
     const aiOpt = {
@@ -11720,20 +12385,20 @@ var AxBootstrapFewShot = class {
           continue;
         }
         const exList = examples.filter((e) => e !== ex);
-        this.program.setExamples(exList);
-        const aiService = this.teacherAI || this.ai;
+        program.setExamples(exList);
+        const aiService = this.getTeacherOrStudentAI();
         this.stats.totalCalls++;
         let res;
         let error;
         try {
-          res = await this.program.forward(aiService, ex, aiOpt);
+          res = await program.forward(aiService, ex, aiOpt);
           if (this.costMonitoring) {
             this.stats.estimatedTokenUsage += JSON.stringify(ex).length / 4 + JSON.stringify(res).length / 4;
           }
-          const score = metricFn({ prediction: res, example: ex });
+          const score = await metricFn({ prediction: res, example: ex });
           const success = score >= 0.5;
           if (success) {
-            this.traces = [...this.traces, ...this.program.getTraces()];
+            this.traces = [...this.traces, ...program.getTraces()];
             this.stats.successfulDemos++;
           }
         } catch (err) {
@@ -11784,13 +12449,15 @@ var AxBootstrapFewShot = class {
       if (!this.stats.earlyStopping) {
         this.stats.earlyStopping = {
           bestScoreRound: improvement > 0 ? roundIndex : 0,
-          patienceExhausted: false
+          patienceExhausted: false,
+          reason: "No improvement detected"
         };
       } 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;
+        this.stats.earlyStopping.reason = `No improvement for ${this.earlyStoppingPatience} rounds`;
         if (this.verboseMode || this.debugMode) {
           console.log(
             `
@@ -11801,37 +12468,38 @@ Early stopping triggered after ${roundIndex + 1} rounds. No improvement for ${th
       }
     }
   }
-  async compile(metricFn, options) {
-    const compileOptions = options;
-    const maxRounds = compileOptions?.maxRounds ?? this.maxRounds;
+  async compile(program, metricFn, options) {
+    const maxRounds = options?.maxIterations ?? this.maxRounds;
     this.traces = [];
-    this.stats = {
-      totalCalls: 0,
-      successfulDemos: 0,
-      estimatedTokenUsage: 0,
-      earlyStopped: false
-    };
+    this.reset();
     for (let i = 0; i < maxRounds; i++) {
-      await this.compileRound(i, metricFn, compileOptions);
+      await this.compileRound(program, i, metricFn, options);
       if (this.stats.earlyStopped) {
         break;
       }
     }
     if (this.traces.length === 0) {
       throw new Error(
-        "No demonstrations found. Either provider more examples or improve the existing ones."
+        "No demonstrations found. Either provide more examples or improve the existing ones."
       );
     }
     const demos = groupTracesByKeys(this.traces);
+    let bestScore = 0;
+    if (this.traces.length > 0) {
+      bestScore = this.stats.successfulDemos / Math.max(1, this.stats.totalCalls);
+    }
     return {
       demos,
-      stats: this.stats
+      stats: this.stats,
+      bestScore,
+      finalConfiguration: {
+        maxRounds: this.maxRounds,
+        maxDemos: this.maxDemos,
+        batchSize: this.batchSize,
+        successRate: bestScore
+      }
     };
   }
-  // Get optimization statistics
-  getStats() {
-    return this.stats;
-  }
 };
 function groupTracesByKeys(programTraces) {
   const groupedTraces = /* @__PURE__ */ new Map();
@@ -11846,9 +12514,12 @@ function groupTracesByKeys(programTraces) {
     }
   }
   const programDemosArray = [];
-  for (const [programId, traces] of groupedTraces.entries()) {
-    programDemosArray.push({ traces, programId });
-  }
+  groupedTraces.forEach((traces, programId) => {
+    programDemosArray.push({
+      traces,
+      programId
+    });
+  });
   return programDemosArray;
 }
 var randomSample = (array, n) => {
@@ -11867,10 +12538,8 @@ var randomSample = (array, n) => {
 };
 
 // dsp/optimizers/miproV2.ts
-var AxMiPRO = class {
-  ai;
-  program;
-  examples;
+var AxMiPRO = class extends AxBaseOptimizer {
+  // MiPRO-specific options
   maxBootstrappedDemos;
   maxLabeledDemos;
   numCandidates;
@@ -11884,52 +12553,35 @@ var AxMiPRO = class {
   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
-      }
-    });
+  bayesianOptimization;
+  acquisitionFunction;
+  explorationWeight;
+  constructor(args) {
+    super(args);
+    const options = args.options || {};
+    this.numCandidates = options.numCandidates ?? 5;
+    this.initTemperature = options.initTemperature ?? 0.7;
+    this.maxBootstrappedDemos = options.maxBootstrappedDemos ?? 3;
+    this.maxLabeledDemos = options.maxLabeledDemos ?? 4;
+    this.numTrials = options.numTrials ?? 30;
+    this.minibatch = options.minibatch ?? true;
+    this.minibatchSize = options.minibatchSize ?? 25;
+    this.minibatchFullEvalSteps = options.minibatchFullEvalSteps ?? 10;
+    this.programAwareProposer = options.programAwareProposer ?? true;
+    this.dataAwareProposer = options.dataAwareProposer ?? true;
+    this.viewDataBatchSize = options.viewDataBatchSize ?? 10;
+    this.tipAwareProposer = options.tipAwareProposer ?? true;
+    this.fewshotAwareProposer = options.fewshotAwareProposer ?? true;
+    this.verbose = options.verbose ?? false;
+    this.earlyStoppingTrials = options.earlyStoppingTrials ?? 5;
+    this.minImprovementThreshold = options.minImprovementThreshold ?? 0.01;
+    this.bayesianOptimization = options.bayesianOptimization ?? false;
+    this.acquisitionFunction = options.acquisitionFunction ?? "expected_improvement";
+    this.explorationWeight = options.explorationWeight ?? 0.1;
+    this.stats.convergenceInfo.convergenceThreshold = this.minImprovementThreshold;
   }
   /**
    * Configures the optimizer for light, medium, or heavy optimization
@@ -11973,123 +12625,60 @@ var AxMiPRO = class {
     ];
   }
   /**
-   * Generates instruction candidates for each predictor in the program
+   * Generates instruction candidates using the teacher model if available
+   * @param options Optional compile options that may override teacher AI
    * @returns Array of generated instruction candidates
    */
-  async proposeInstructionCandidates() {
+  async proposeInstructionCandidates(options) {
     const instructions = [];
-    let programContext = "";
-    if (this.programAwareProposer) {
-      programContext = await this.generateProgramSummary();
-    }
-    let dataContext = "";
-    if (this.dataAwareProposer) {
-      dataContext = await this.generateDataSummary();
-    }
+    const aiToUse = this.getTeacherOrStudentAI(options);
     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
+        candidateIndex: i,
+        ai: aiToUse
       });
       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 "";
+    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 response.results[0]?.content || "";
+    return instruction;
   }
   /**
    * Bootstraps few-shot examples for the program
    */
-  async bootstrapFewShotExamples(metricFn) {
+  async bootstrapFewShotExamples(program, metricFn) {
     if (this.verbose) {
       console.log("Bootstrapping few-shot examples...");
     }
-    const result = await this.bootstrapper.compile(metricFn, {
+    const bootstrapper = new AxBootstrapFewShot({
+      studentAI: this.studentAI,
+      examples: this.examples,
+      options: {
+        maxDemos: this.maxBootstrappedDemos,
+        maxRounds: 3,
+        verboseMode: this.verbose
+      }
+    });
+    const result = await bootstrapper.compile(program, metricFn, {
       maxDemos: this.maxBootstrappedDemos
     });
     return result.demos || [];
@@ -12113,109 +12702,98 @@ ${dataContext}
     return selectedExamples;
   }
   /**
-   * Runs Bayesian optimization to find the best combination of few-shot examples and instructions
+   * Runs 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 = {
+  async runOptimization(program, bootstrappedDemos, labeledExamples, instructions, valset, metricFn, options) {
+    let bestConfig = {
       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] || "";
+    let bestScore = 0;
+    let stagnationRounds = 0;
+    const scoreHistory = [];
+    let startRound = 0;
+    if (this.resumeFromCheckpoint) {
+      const checkpoint = await this.loadCheckpoint(
+        this.resumeFromCheckpoint,
+        options
+      );
+      if (checkpoint && checkpoint.optimizerType === "MiPRO") {
+        if (this.verbose || options?.verbose) {
+          console.log(
+            `Resuming from checkpoint at round ${checkpoint.currentRound}`
+          );
+        }
+        this.restoreFromCheckpoint(checkpoint);
+        startRound = checkpoint.currentRound;
+        bestScore = checkpoint.bestScore;
+        bestConfig = checkpoint.bestConfiguration || bestConfig;
+        stagnationRounds = checkpoint.stats.convergenceInfo?.stagnationRounds || 0;
+      }
+    }
+    for (let i = startRound; i < this.numTrials; i++) {
       const config = {
-        instruction: instructionValue,
-        bootstrappedDemos: Math.floor(
-          Math.random() * (bootstrappedDemos.length + 1)
+        instruction: instructions[i % instructions.length] || instructions[0] || "",
+        bootstrappedDemos: Math.min(
+          Math.floor(Math.random() * (bootstrappedDemos.length + 1)),
+          this.maxBootstrappedDemos
         ),
-        labeledExamples: Math.floor(
-          Math.random() * (labeledExamples.length + 1)
+        labeledExamples: Math.min(
+          Math.floor(Math.random() * (labeledExamples.length + 1)),
+          this.maxLabeledDemos
         )
       };
-      configs.push(config);
-    }
-    for (let i = 0; i < configs.length; i++) {
-      const config = configs[i];
-      if (!config) continue;
       const score = await this.evaluateConfig(
+        program,
         config,
         bootstrappedDemos,
         labeledExamples,
         valset,
-        metricFn,
-        i
+        metricFn
       );
-      evaluatedConfigs.push({ config, score });
-      if (score > bestScore) {
+      scoreHistory.push(score);
+      const improvement = score - bestScore;
+      if (improvement > this.minImprovementThreshold) {
         bestScore = score;
         bestConfig = config;
-        if (this.verbose) {
-          console.log(
-            `New best configuration found with score ${bestScore} (exploration phase)`
-          );
-        }
+        stagnationRounds = 0;
+      } else {
+        stagnationRounds++;
       }
-      updateProgressBar(
+      await this.updateOptimizationProgress(
         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
+        score,
+        config,
+        "MiPRO",
+        this.getConfiguration(),
+        bestScore,
+        bestConfig,
+        {
+          stagnationRounds,
+          bootstrappedDemos: bootstrappedDemos.length,
+          labeledExamples: labeledExamples.length,
+          instructions: instructions.length
+        },
+        options
       );
-      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;
+      if (this.onProgress) {
+        this.onProgress({
+          round: i + 1,
+          totalRounds: this.numTrials,
+          currentScore: score,
+          bestScore,
+          tokensUsed: this.stats.resourceUsage.totalTokens,
+          timeElapsed: Date.now(),
+          successfulExamples: this.stats.successfulDemos,
+          totalExamples: this.examples.length,
+          currentConfiguration: config,
+          convergenceInfo: {
+            improvement,
+            stagnationRounds,
+            isConverging: stagnationRounds < this.earlyStoppingTrials
           }
-        } else {
-          lastBestScore = bestScore;
-          trialsWithoutImprovement = 0;
-        }
+        });
       }
       updateProgressBar(
         i + 1,
@@ -12225,243 +12803,91 @@ ${dataContext}
         "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 (this.checkCostLimits()) {
+        this.triggerEarlyStopping("Cost limit reached", i + 1);
+        break;
       }
-    }
-    if (!bestConfig) {
-      if (this.verbose) {
-        console.warn(
-          "Optimization failed to find any valid configurations, using default fallback configuration"
+      if (stagnationRounds >= this.earlyStoppingTrials) {
+        this.triggerEarlyStopping(
+          `No improvement for ${this.earlyStoppingTrials} trials`,
+          i - stagnationRounds + 1
         );
+        break;
       }
-      bestConfig = defaultConfig;
-      try {
-        bestScore = await this.evaluateConfig(
-          bestConfig,
-          bootstrappedDemos,
-          labeledExamples,
-          valset,
-          metricFn,
-          this.numTrials - 1
+      if (this.checkTargetScore(bestScore)) {
+        this.triggerEarlyStopping(
+          `Target score ${this.targetScore} reached`,
+          i + 1
         );
-      } catch (err) {
-        if (this.verbose) {
-          console.error("Error evaluating default configuration:", err);
-        }
-        bestScore = 0;
+        break;
       }
     }
+    this.stats.convergenceInfo.stagnationRounds = stagnationRounds;
+    this.stats.convergenceInfo.finalImprovement = scoreHistory.length > 1 ? bestScore - scoreHistory[0] : 0;
+    this.stats.convergenceInfo.converged = stagnationRounds < this.earlyStoppingTrials;
     return { bestConfig, bestScore };
   }
-  /**
-   * Evaluates a configuration on the validation set
-   */
-  async evaluateConfig(config, bootstrappedDemos, labeledExamples, valset, metricFn, trialIndex) {
+  async evaluateConfig(program, config, bootstrappedDemos, labeledExamples, valset, metricFn) {
+    const testProgram = { ...program };
     this.applyConfigToProgram(
-      this.program,
+      testProgram,
       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 sumOfScores = 0;
+    let totalScore = 0;
+    let count = 0;
+    const evalSet = valset.slice(0, Math.min(5, valset.length));
     for (const example of evalSet) {
       try {
-        const prediction = await this.program.forward(this.ai, example);
-        const score = metricFn({ prediction, example });
-        sumOfScores += score;
-      } catch (err) {
-        if (this.verbose) {
-          console.error("Error evaluating example:", err);
-        }
-      }
-    }
-    if (evalSet.length === 0) return 0;
-    return sumOfScores / 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 sumOfScores = 0;
-    for (const example of valset) {
-      try {
-        const prediction = await this.program.forward(this.ai, example);
-        const score = metricFn({ prediction, example });
-        sumOfScores += score;
-      } catch (err) {
-        if (this.verbose) {
-          console.error("Error evaluating example:", err);
-        }
+        const prediction = await testProgram.forward(
+          this.studentAI,
+          example
+        );
+        const score = await metricFn({ prediction, example });
+        totalScore += score;
+        count++;
+        this.stats.totalCalls++;
+      } catch {
+        continue;
       }
     }
-    if (valset.length === 0) return 0;
-    return sumOfScores / valset.length;
+    return count > 0 ? totalScore / count : 0;
   }
-  /**
-   * 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) {
+    if (program.setInstruction) {
+      program.setInstruction(config.instruction);
+    }
+    if (config.bootstrappedDemos > 0 && program.setDemos) {
       program.setDemos(bootstrappedDemos.slice(0, config.bootstrappedDemos));
     }
-    if (config.labeledExamples > 0) {
+    if (config.labeledExamples > 0 && program.setExamples) {
       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 optimization result
    */
-  async compile(metricFn, options) {
+  async compile(program, metricFn, options) {
+    const startTime = Date.now();
+    this.setupRandomSeed();
     const miproOptions = options;
     if (miproOptions?.auto) {
       this.configureAuto(miproOptions.auto);
     }
-    const trainset = this.examples;
-    const valset = miproOptions?.valset || this.examples.slice(0, Math.floor(this.examples.length * 0.8));
-    if (this.verbose) {
+    const valset = this.getValidationSet(options) || (miproOptions?.valset ?? this.examples.slice(0, Math.floor(this.examples.length * 0.2)));
+    if (this.verbose || options?.verbose) {
       console.log(`Starting MIPROv2 optimization with ${this.numTrials} trials`);
       console.log(
-        `Using ${trainset.length} examples for training and ${valset.length} for validation`
+        `Using ${this.examples.length} examples for training and ${valset.length} for validation`
       );
-    }
-    if (miproOptions?.teacher) {
-      if (this.verbose) {
-        console.log("Using provided teacher to assist with bootstrapping");
+      if (this.teacherAI) {
+        console.log("Using separate teacher model for instruction generation");
       }
-      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);
+      bootstrappedDemos = await this.bootstrapFewShotExamples(program, metricFn);
       if (this.verbose) {
         console.log(
           `Generated ${bootstrappedDemos.length} bootstrapped demonstrations`
@@ -12477,38 +12903,191 @@ ${dataContext}
         );
       }
     }
-    const instructions = await this.proposeInstructionCandidates();
+    const instructions = await this.proposeInstructionCandidates(options);
     if (this.verbose) {
       console.log(`Generated ${instructions.length} instruction candidates`);
+      if (this.hasTeacherAI(options)) {
+        console.log("Using teacher AI for instruction generation");
+      }
     }
-    const { bestConfig, bestScore } = await this.runBayesianOptimization(
+    const { bestConfig, bestScore } = await this.runOptimization(
+      program,
       bootstrappedDemos,
       labeledExamples,
       instructions,
       valset,
-      metricFn
+      metricFn,
+      options
     );
-    if (this.verbose) {
+    if (this.verbose || options?.verbose) {
       console.log(`Optimization complete. Best score: ${bestScore}`);
       console.log(`Best configuration: ${JSON.stringify(bestConfig)}`);
     }
-    this.applyConfigToProgram(
-      this.program,
+    if (this.checkTargetScore(bestScore)) {
+      this.triggerEarlyStopping(
+        `Target score ${this.targetScore} reached with score ${bestScore}`,
+        this.numTrials
+      );
+    }
+    let signature;
+    if ("getSignature" in program && typeof program.getSignature === "function") {
+      signature = program.getSignature();
+    } else {
+      signature = "input -> output";
+    }
+    const optimizedGen = new AxGen(signature);
+    this.applyConfigToAxGen(
+      optimizedGen,
       bestConfig,
       bootstrappedDemos,
       labeledExamples
     );
+    this.updateResourceUsage(startTime);
+    this.stats.convergenceInfo.converged = true;
+    this.stats.convergenceInfo.finalImprovement = bestScore;
+    await this.saveFinalCheckpoint(
+      "MiPRO",
+      this.getConfiguration(),
+      bestScore,
+      bestConfig,
+      {
+        bootstrappedDemos: bootstrappedDemos.length,
+        labeledExamples: labeledExamples.length,
+        instructions: instructions.length,
+        optimizedGen: !!optimizedGen
+      },
+      options
+    );
     return {
-      program: this.program,
-      demos: bootstrappedDemos
+      demos: bootstrappedDemos,
+      stats: this.stats,
+      bestScore,
+      optimizedGen,
+      finalConfiguration: {
+        instruction: bestConfig.instruction,
+        bootstrappedDemos: bestConfig.bootstrappedDemos,
+        labeledExamples: bestConfig.labeledExamples,
+        numCandidates: this.numCandidates,
+        numTrials: this.numTrials
+      }
     };
   }
   /**
-   * Get optimization statistics from the internal bootstrapper
-   * @returns Optimization statistics or undefined if not available
+   * Applies a configuration to an AxGen instance
    */
-  getStats() {
-    return this.bootstrapper.getStats();
+  applyConfigToAxGen(axgen, config, bootstrappedDemos, labeledExamples) {
+    if ("setInstruction" in axgen && typeof axgen.setInstruction === "function") {
+      axgen.setInstruction(config.instruction);
+    }
+    if (config.bootstrappedDemos > 0) {
+      axgen.setDemos(bootstrappedDemos.slice(0, config.bootstrappedDemos));
+    }
+    if (config.labeledExamples > 0) {
+      axgen.setExamples(
+        labeledExamples.slice(
+          0,
+          config.labeledExamples
+        )
+      );
+    }
+  }
+  /**
+   * Get optimizer-specific configuration
+   * @returns Current optimizer configuration
+   */
+  getConfiguration() {
+    return {
+      numCandidates: this.numCandidates,
+      initTemperature: this.initTemperature,
+      maxBootstrappedDemos: this.maxBootstrappedDemos,
+      maxLabeledDemos: this.maxLabeledDemos,
+      numTrials: this.numTrials,
+      minibatch: this.minibatch,
+      minibatchSize: this.minibatchSize,
+      minibatchFullEvalSteps: this.minibatchFullEvalSteps,
+      programAwareProposer: this.programAwareProposer,
+      dataAwareProposer: this.dataAwareProposer,
+      tipAwareProposer: this.tipAwareProposer,
+      fewshotAwareProposer: this.fewshotAwareProposer,
+      earlyStoppingTrials: this.earlyStoppingTrials,
+      minImprovementThreshold: this.minImprovementThreshold,
+      bayesianOptimization: this.bayesianOptimization,
+      acquisitionFunction: this.acquisitionFunction,
+      explorationWeight: this.explorationWeight
+    };
+  }
+  /**
+   * Update optimizer configuration
+   * @param config New configuration to merge with existing
+   */
+  updateConfiguration(config) {
+    if (config.numCandidates !== void 0) {
+      this.numCandidates = config.numCandidates;
+    }
+    if (config.initTemperature !== void 0) {
+      this.initTemperature = config.initTemperature;
+    }
+    if (config.maxBootstrappedDemos !== void 0) {
+      this.maxBootstrappedDemos = config.maxBootstrappedDemos;
+    }
+    if (config.maxLabeledDemos !== void 0) {
+      this.maxLabeledDemos = config.maxLabeledDemos;
+    }
+    if (config.numTrials !== void 0) {
+      this.numTrials = config.numTrials;
+    }
+    if (config.minibatch !== void 0) {
+      this.minibatch = config.minibatch;
+    }
+    if (config.minibatchSize !== void 0) {
+      this.minibatchSize = config.minibatchSize;
+    }
+    if (config.earlyStoppingTrials !== void 0) {
+      this.earlyStoppingTrials = config.earlyStoppingTrials;
+    }
+    if (config.minImprovementThreshold !== void 0) {
+      this.minImprovementThreshold = config.minImprovementThreshold;
+    }
+    if (config.verbose !== void 0) {
+      this.verbose = config.verbose;
+    }
+  }
+  /**
+   * Reset optimizer state for reuse with different programs
+   */
+  reset() {
+    super.reset();
+    this.stats.convergenceInfo.convergenceThreshold = this.minImprovementThreshold;
+  }
+  /**
+   * Validate that the optimizer can handle the given program
+   * @param program Program to validate
+   * @returns Validation result with any issues found
+   */
+  validateProgram(program) {
+    const result = super.validateProgram(program);
+    if (this.examples.length < this.maxBootstrappedDemos + this.maxLabeledDemos) {
+      result.issues.push(
+        `Not enough examples: need at least ${this.maxBootstrappedDemos + this.maxLabeledDemos}, got ${this.examples.length}`
+      );
+      result.suggestions.push(
+        "Reduce maxBootstrappedDemos or maxLabeledDemos, or provide more examples"
+      );
+    }
+    const valSetSize = this.getValidationSet().length;
+    if (valSetSize < 5) {
+      result.issues.push(
+        "Validation set too small for reliable MiPRO optimization"
+      );
+      result.suggestions.push(
+        "Provide more examples or a larger validation set"
+      );
+    }
+    return {
+      isValid: result.issues.length === 0,
+      issues: result.issues,
+      suggestions: result.suggestions
+    };
   }
 };
 
@@ -12755,7 +13334,7 @@ var AxTestPrompt = class {
         throw new Error("Invalid example");
       }
       const res = await this.program.forward(this.ai, ex);
-      const score = metricFn({ prediction: res, example: ex });
+      const score = await metricFn({ prediction: res, example: ex });
       sumOfScores += score;
       const et = (/* @__PURE__ */ new Date()).getTime() - st;
       updateProgressBar(i, total, sumOfScores, et, "Testing Prompt", 30);
@@ -14789,7 +15368,6 @@ var AxRAG = class extends AxChainOfThought {
     );
     this.genQuery = new AxGen(qsig);
     this.queryFn = queryFn;
-    this.register(this.genQuery);
   }
   async forward(ai, values, options) {
     let question;
@@ -14867,6 +15445,7 @@ var AxRAG = class extends AxChainOfThought {
   AxAssertionError,
   AxBalancer,
   AxBaseAI,
+  AxBaseOptimizer,
   AxBootstrapFewShot,
   AxChainOfThought,
   AxDB,
@@ -14876,6 +15455,7 @@ var AxRAG = class extends AxChainOfThought {
   AxDBMemory,
   AxDBPinecone,
   AxDBWeaviate,
+  AxDefaultCostTracker,
   AxDefaultQueryRewriter,
   AxDefaultResultReranker,
   AxDockerSession,