@ax-llm/ax 12.0.22 → 12.0.24

package/index.cjs

@@ -117,7 +117,6 @@ __export(index_exports, {
   AxMockAIService: () => AxMockAIService,
   AxMultiServiceRouter: () => AxMultiServiceRouter,
   AxProgram: () => AxProgram,
-  AxProgramWithSignature: () => AxProgramWithSignature,
   AxPromptTemplate: () => AxPromptTemplate,
   AxRAG: () => AxRAG,
   AxRateLimiterTokenUsage: () => AxRateLimiterTokenUsage,
@@ -4987,7 +4986,8 @@ var AxAIGoogleGemini = class extends AxBaseAI {
       } else {
         path = "publishers/google";
       }
-      apiURL = `https://${region}-aiplatform.googleapis.com/v1/projects/${projectId}/locations/${region}/${path}`;
+      const tld = region === "global" ? "aiplatform" : `${region}-aiplatform`;
+      apiURL = `https://${tld}.googleapis.com/v1/projects/${projectId}/locations/${region}/${path}`;
       if (apiKey) {
         headers = async () => ({ Authorization: `Bearer ${apiKey}` });
       } else {
@@ -11277,7 +11277,7 @@ var AxInstanceRegistry = class {
 };
 
 // dsp/program.ts
-var AxProgramWithSignature = class {
+var AxProgram = class {
   signature;
   sigHash;
   examples;
@@ -11285,6 +11285,7 @@ var AxProgramWithSignature = class {
   demos;
   trace;
   usage = [];
+  traceLabel;
   key;
   children;
   constructor(signature, options) {
@@ -11292,6 +11293,9 @@ var AxProgramWithSignature = class {
     if (options?.description) {
       this.signature.setDescription(options.description);
     }
+    if (options?.traceLabel) {
+      this.traceLabel = options.traceLabel;
+    }
     this.signature.validate();
     this.sigHash = this.signature?.hash();
     this.children = new AxInstanceRegistry();
@@ -11397,85 +11401,6 @@ var AxProgramWithSignature = class {
     }
   }
 };
-var AxProgram = class {
-  trace;
-  usage = [];
-  key;
-  children;
-  constructor() {
-    this.children = new AxInstanceRegistry();
-    this.key = { id: this.constructor.name };
-  }
-  register(prog) {
-    if (this.key) {
-      prog.setParentId(this.key.id);
-    }
-    this.children.register(prog);
-  }
-  async forward(_ai, _values, _options) {
-    throw new Error("forward() not implemented");
-  }
-  // biome-ignore lint/correctness/useYield: just a placeholder
-  async *streamingForward(_ai, _values, _options) {
-    throw new Error("streamingForward() not implemented");
-  }
-  setId(id) {
-    this.key = { id, custom: true };
-    for (const child of Array.from(this.children)) {
-      child?.setParentId(id);
-    }
-  }
-  setParentId(parentId) {
-    if (!this.key.custom) {
-      this.key.id = [parentId, this.key.id].join("/");
-    }
-  }
-  setExamples(examples, options) {
-    if (!("programId" in examples)) {
-      return;
-    }
-    for (const child of Array.from(this.children)) {
-      child?.setExamples(examples, options);
-    }
-  }
-  getTraces() {
-    let traces = [];
-    if (this.trace) {
-      traces.push({ trace: this.trace, programId: this.key.id });
-    }
-    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 Array.from(this.children)) {
-      const cu = child?.getUsage();
-      usage = [...usage, ...cu ?? []];
-    }
-    return mergeProgramUsage(usage);
-  }
-  resetUsage() {
-    this.usage = [];
-    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`
-      );
-    }
-    for (const child of Array.from(this.children)) {
-      child?.setDemos(demos);
-    }
-  }
-};
 
 // dsp/samples.ts
 function checkForFunctionCalls(mem, sessionId) {
@@ -11590,7 +11515,7 @@ ${errors}`, {
 }
 
 // dsp/generate.ts
-var AxGen = class extends AxProgramWithSignature {
+var AxGen = class extends AxProgram {
   promptTemplate;
   asserts;
   streamingAsserts;
@@ -11601,7 +11526,10 @@ var AxGen = class extends AxProgramWithSignature {
   excludeContentFromTrace = false;
   thoughtFieldName;
   constructor(signature, options) {
-    super(signature, { description: options?.description });
+    super(signature, {
+      description: options?.description,
+      traceLabel: options?.traceLabel
+    });
     this.options = options;
     this.thoughtFieldName = options?.thoughtFieldName ?? "thought";
     const promptTemplateOptions = {
@@ -12045,8 +11973,8 @@ var AxGen = class extends AxProgramWithSignature {
       ...options?.maxSteps ? { max_steps: options.maxSteps } : {},
       ...options?.maxRetries ? { max_retries: options.maxRetries } : {}
     };
-    const traceLabel = options.traceLabel ?? this.options?.traceLabel;
-    const spanName = traceLabel ? `${traceLabel} (AxGen)` : "AxGen";
+    const traceLabel = this.traceLabel && options.traceLabel ? `${this.traceLabel} > ${options.traceLabel}` : options.traceLabel ?? this.traceLabel;
+    const spanName = traceLabel ? `AxGen > ${traceLabel}` : "AxGen";
     const span = tracer.startSpan(spanName, {
       kind: import_api22.SpanKind.SERVER,
       attributes
@@ -15182,15 +15110,186 @@ var AxDockerSession = class {
 };
 
 // flow/flow.ts
-var AxFlow = class extends AxProgramWithSignature {
+var AxFlowDependencyAnalyzer = class {
+  /**
+   * Analyzes a mapping function to determine which state fields it depends on
+   */
+  analyzeMappingDependencies(mapping, _nodeName) {
+    const dependencies = [];
+    const source = mapping.toString();
+    const stateAccessMatches = Array.from(source.matchAll(/state\.(\w+)/g));
+    for (const match of stateAccessMatches) {
+      if (match[1] && !dependencies.includes(match[1])) {
+        dependencies.push(match[1]);
+      }
+    }
+    if (dependencies.length === 0) {
+      try {
+        const tracker = this.createDependencyTracker(dependencies);
+        mapping(tracker);
+      } catch {
+      }
+    }
+    return dependencies;
+  }
+  createDependencyTracker(dependencies) {
+    return new Proxy(
+      {},
+      {
+        get(target, prop) {
+          if (typeof prop === "string" && !dependencies.includes(prop)) {
+            dependencies.push(prop);
+          }
+          return new Proxy(
+            {},
+            {
+              get: () => void 0
+            }
+          );
+        }
+      }
+    );
+  }
+};
+var AxFlowExecutionPlanner = class {
+  steps = [];
+  parallelGroups = [];
+  analyzer = new AxFlowDependencyAnalyzer();
+  initialFields = /* @__PURE__ */ new Set();
+  /**
+   * Adds an execution step to the plan
+   */
+  addExecutionStep(stepFunction, nodeName, mapping) {
+    let dependencies = [];
+    let produces = [];
+    let type = "other";
+    if (nodeName && mapping) {
+      type = "execute";
+      dependencies = this.analyzer.analyzeMappingDependencies(mapping, nodeName);
+      produces = [`${nodeName}Result`];
+    } else if (stepFunction.toString().includes("transform(")) {
+      type = "map";
+      dependencies = this.getAllProducedFields();
+    }
+    const step = {
+      type,
+      nodeName,
+      dependencies,
+      produces,
+      stepFunction,
+      stepIndex: this.steps.length
+    };
+    this.steps.push(step);
+  }
+  /**
+   * Sets the initial fields and rebuilds parallel groups
+   */
+  setInitialFields(fields) {
+    this.initialFields = new Set(fields);
+    this.rebuildParallelGroups();
+  }
+  /**
+   * Rebuilds the parallel execution groups based on dependencies
+   */
+  rebuildParallelGroups() {
+    this.parallelGroups = [];
+    const processedSteps = /* @__PURE__ */ new Set();
+    const availableFields = new Set(this.initialFields);
+    let currentLevel = 0;
+    while (processedSteps.size < this.steps.length) {
+      const currentLevelSteps = [];
+      for (const step of this.steps) {
+        if (processedSteps.has(step.stepIndex)) continue;
+        const canRun = step.dependencies.length === 0 || step.dependencies.every((dep) => availableFields.has(dep));
+        if (canRun) {
+          currentLevelSteps.push(step);
+          processedSteps.add(step.stepIndex);
+        }
+      }
+      if (currentLevelSteps.length > 0) {
+        for (const step of currentLevelSteps) {
+          step.produces.forEach((field) => availableFields.add(field));
+        }
+        this.parallelGroups.push({
+          level: currentLevel,
+          steps: currentLevelSteps
+        });
+        currentLevel++;
+      } else {
+        break;
+      }
+    }
+  }
+  /**
+   * Gets all fields produced by previous steps
+   */
+  getAllProducedFields() {
+    const fields = [];
+    for (const step of this.steps) {
+      fields.push(...step.produces);
+    }
+    return fields;
+  }
+  /**
+   * Creates optimized execution function
+   */
+  createOptimizedExecution() {
+    const optimizedSteps = [];
+    for (const group of this.parallelGroups) {
+      if (group.steps.length === 1) {
+        const step = group.steps[0];
+        if (step) {
+          optimizedSteps.push(step.stepFunction);
+        }
+      } else if (group.steps.length > 1) {
+        const parallelStep = async (state, context3) => {
+          const promises = group.steps.map(
+            (step) => step.stepFunction(state, context3)
+          );
+          const results = await Promise.all(promises);
+          let mergedState = state;
+          for (const result of results) {
+            mergedState = { ...mergedState, ...result };
+          }
+          return mergedState;
+        };
+        optimizedSteps.push(parallelStep);
+      }
+    }
+    return optimizedSteps;
+  }
+  /**
+   * Gets execution plan info for debugging
+   */
+  getExecutionPlan() {
+    return {
+      totalSteps: this.steps.length,
+      parallelGroups: this.parallelGroups.length,
+      maxParallelism: Math.max(
+        ...this.parallelGroups.map((g) => g.steps.length),
+        0
+      ),
+      steps: this.steps,
+      groups: this.parallelGroups
+    };
+  }
+};
+var AxFlow = class extends AxProgram {
   nodes = /* @__PURE__ */ new Map();
   flowDefinition = [];
   nodeGenerators = /* @__PURE__ */ new Map();
   loopStack = [];
   stepLabels = /* @__PURE__ */ new Map();
   branchContext = null;
-  constructor(signature = "userInput:string -> flowOutput:string") {
+  // Automatic parallelization components
+  autoParallelConfig;
+  executionPlanner = new AxFlowExecutionPlanner();
+  constructor(signature = "userInput:string -> flowOutput:string", options) {
     super(signature);
+    this.autoParallelConfig = {
+      enabled: options?.autoParallel !== false
+      // Default to true
+    };
   }
   // Implementation
   node(name, signatureOrAxGenOrClass, options) {
@@ -15209,7 +15308,7 @@ var AxFlow = class extends AxProgramWithSignature {
         outputs: {}
       });
       this.nodeGenerators.set(name, new AxGen(signatureOrAxGenOrClass, options));
-    } else if (typeof signatureOrAxGenOrClass === "function" && signatureOrAxGenOrClass.prototype instanceof AxProgramWithSignature) {
+    } else if (typeof signatureOrAxGenOrClass === "function" && signatureOrAxGenOrClass.prototype instanceof AxProgram) {
       this.nodes.set(name, {
         inputs: {},
         outputs: {}
@@ -15230,7 +15329,7 @@ var AxFlow = class extends AxProgramWithSignature {
       this.nodeGenerators.set(name, new AxGen(signature, options));
     } else {
       throw new Error(
-        `Invalid second argument for node '${name}': expected string, AxSignature, AxGen instance, or class extending AxProgramWithSignature`
+        `Invalid second argument for node '${name}': expected string, AxSignature, AxGen instance, or class extending AxProgram`
       );
     }
     return this;
@@ -15265,6 +15364,9 @@ var AxFlow = class extends AxProgramWithSignature {
       );
     } else {
       this.flowDefinition.push(step);
+      if (this.autoParallelConfig.enabled) {
+        this.executionPlanner.addExecutionStep(step);
+      }
     }
     return this;
   }
@@ -15327,7 +15429,11 @@ var AxFlow = class extends AxProgramWithSignature {
       const ai = dynamicContext?.ai ?? context3.mainAi;
       const options = dynamicContext?.options ?? context3.mainOptions;
       const nodeInputs = mapping(state);
-      const result = await nodeProgram.forward(ai, nodeInputs, options);
+      const traceLabel = options?.traceLabel ? `Node:${nodeName} (${options.traceLabel})` : `Node:${nodeName}`;
+      const result = await nodeProgram.forward(ai, nodeInputs, {
+        ...options,
+        traceLabel
+      });
       return {
         ...state,
         [`${nodeName}Result`]: result
@@ -15344,6 +15450,9 @@ var AxFlow = class extends AxProgramWithSignature {
       );
     } else {
       this.flowDefinition.push(step);
+      if (this.autoParallelConfig.enabled) {
+        this.executionPlanner.addExecutionStep(step, nodeName, mapping);
+      }
     }
     return this;
   }
@@ -15640,7 +15749,7 @@ var AxFlow = class extends AxProgramWithSignature {
    *
    * @param ai - The AI service to use as the default for all steps
    * @param values - The input values for the flow
-   * @param options - Optional forward options to use as defaults
+   * @param options - Optional forward options to use as defaults (includes autoParallel override)
    * @returns Promise that resolves to the final output
    */
   async forward(ai, values, options) {
@@ -15649,11 +15758,36 @@ var AxFlow = class extends AxProgramWithSignature {
       mainAi: ai,
       mainOptions: options
     };
-    for (const step of this.flowDefinition) {
-      state = await step(state, context3);
+    const useAutoParallel = options?.autoParallel !== false && this.autoParallelConfig.enabled;
+    if (useAutoParallel) {
+      this.executionPlanner.setInitialFields(Object.keys(values));
+      const optimizedSteps = this.executionPlanner.createOptimizedExecution();
+      for (const step of optimizedSteps) {
+        state = await step(state, context3);
+      }
+    } else {
+      for (const step of this.flowDefinition) {
+        state = await step(state, context3);
+      }
     }
     return state;
   }
+  /**
+   * Gets execution plan information for debugging automatic parallelization
+   *
+   * @returns Object with execution plan details
+   */
+  getExecutionPlan() {
+    const planInfo = this.executionPlanner.getExecutionPlan();
+    return {
+      totalSteps: planInfo.totalSteps,
+      parallelGroups: planInfo.parallelGroups,
+      maxParallelism: planInfo.maxParallelism,
+      autoParallelEnabled: this.autoParallelConfig.enabled,
+      steps: planInfo.steps,
+      groups: planInfo.groups
+    };
+  }
 };
 var AxFlowSubContextImpl = class {
   constructor(nodeGenerators) {
@@ -15669,7 +15803,11 @@ var AxFlowSubContextImpl = class {
       const ai = dynamicContext?.ai ?? context3.mainAi;
       const options = dynamicContext?.options ?? context3.mainOptions;
       const nodeInputs = mapping(state);
-      const result = await nodeProgram.forward(ai, nodeInputs, options);
+      const traceLabel = options?.traceLabel ? `Node:${nodeName} (${options.traceLabel})` : `Node:${nodeName}`;
+      const result = await nodeProgram.forward(ai, nodeInputs, {
+        ...options,
+        traceLabel
+      });
       return {
         ...state,
         [`${nodeName}Result`]: result
@@ -15703,7 +15841,11 @@ var AxFlowTypedSubContextImpl = class {
       const ai = dynamicContext?.ai ?? context3.mainAi;
       const options = dynamicContext?.options ?? context3.mainOptions;
       const nodeInputs = mapping(state);
-      const result = await nodeProgram.forward(ai, nodeInputs, options);
+      const traceLabel = options?.traceLabel ? `Node:${nodeName} (${options.traceLabel})` : `Node:${nodeName}`;
+      const result = await nodeProgram.forward(ai, nodeInputs, {
+        ...options,
+        traceLabel
+      });
       return {
         ...state,
         [`${nodeName}Result`]: result
@@ -16511,10 +16653,7 @@ var AxMiPRO = class extends AxBaseOptimizer {
    * Generates program summary for context-aware instruction generation
    */
   async generateProgramSummary(program, ai) {
-    let signature = "input -> output";
-    if ("getSignature" in program && typeof program.getSignature === "function") {
-      signature = program.getSignature();
-    }
+    const signature = program.getSignature();
     const summaryPrompt = `
 Analyze this language model program and provide a concise summary of its purpose and structure.
 
@@ -19728,7 +19867,6 @@ var AxRAG = class extends AxChainOfThought {
   AxMockAIService,
   AxMultiServiceRouter,
   AxProgram,
-  AxProgramWithSignature,
   AxPromptTemplate,
   AxRAG,
   AxRateLimiterTokenUsage,