@ax-llm/ax 12.0.21 → 12.0.22

package/index.cjs

@@ -15193,18 +15193,31 @@ var AxFlow = class extends AxProgramWithSignature {
     super(signature);
   }
   // Implementation
-  node(name, signatureOrAxGen, options) {
-    if (signatureOrAxGen instanceof AxGen) {
+  node(name, signatureOrAxGenOrClass, options) {
+    if (signatureOrAxGenOrClass instanceof AxGen) {
       this.nodes.set(name, {
         inputs: {},
         outputs: {}
       });
       this.nodeGenerators.set(
         name,
-        signatureOrAxGen
+        signatureOrAxGenOrClass
       );
-    } else {
-      const signature = signatureOrAxGen;
+    } else if (signatureOrAxGenOrClass instanceof AxSignature) {
+      this.nodes.set(name, {
+        inputs: {},
+        outputs: {}
+      });
+      this.nodeGenerators.set(name, new AxGen(signatureOrAxGenOrClass, options));
+    } else if (typeof signatureOrAxGenOrClass === "function" && signatureOrAxGenOrClass.prototype instanceof AxProgramWithSignature) {
+      this.nodes.set(name, {
+        inputs: {},
+        outputs: {}
+      });
+      const programInstance = new signatureOrAxGenOrClass();
+      this.nodeGenerators.set(name, programInstance);
+    } else if (typeof signatureOrAxGenOrClass === "string") {
+      const signature = signatureOrAxGenOrClass;
       if (!signature) {
         throw new Error(
           `Invalid signature for node '${name}': signature cannot be empty`
@@ -15215,11 +15228,15 @@ var AxFlow = class extends AxProgramWithSignature {
         outputs: {}
       });
       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`
+      );
     }
     return this;
   }
-  n(name, signatureOrAxGen, options) {
-    return this.node(name, signatureOrAxGen, options);
+  n(name, signatureOrAxGenOrClass, options) {
+    return this.node(name, signatureOrAxGenOrClass, options);
   }
   /**
    * Applies a synchronous transformation to the state object.
@@ -15302,15 +15319,15 @@ var AxFlow = class extends AxProgramWithSignature {
         `Node '${nodeName}' not found. Make sure to define it with .node() first.`
       );
     }
-    const nodeGenerator = this.nodeGenerators.get(nodeName);
-    if (!nodeGenerator) {
-      throw new Error(`Node generator for '${nodeName}' not found.`);
+    const nodeProgram = this.nodeGenerators.get(nodeName);
+    if (!nodeProgram) {
+      throw new Error(`Node program for '${nodeName}' not found.`);
     }
     const step = async (state, context3) => {
       const ai = dynamicContext?.ai ?? context3.mainAi;
       const options = dynamicContext?.options ?? context3.mainOptions;
       const nodeInputs = mapping(state);
-      const result = await nodeGenerator.forward(ai, nodeInputs, options);
+      const result = await nodeProgram.forward(ai, nodeInputs, options);
       return {
         ...state,
         [`${nodeName}Result`]: result
@@ -15644,15 +15661,15 @@ var AxFlowSubContextImpl = class {
   }
   steps = [];
   execute(nodeName, mapping, dynamicContext) {
-    const nodeGenerator = this.nodeGenerators.get(nodeName);
-    if (!nodeGenerator) {
-      throw new Error(`Node generator for '${nodeName}' not found.`);
+    const nodeProgram = this.nodeGenerators.get(nodeName);
+    if (!nodeProgram) {
+      throw new Error(`Node program for '${nodeName}' not found.`);
     }
     this.steps.push(async (state, context3) => {
       const ai = dynamicContext?.ai ?? context3.mainAi;
       const options = dynamicContext?.options ?? context3.mainOptions;
       const nodeInputs = mapping(state);
-      const result = await nodeGenerator.forward(ai, nodeInputs, options);
+      const result = await nodeProgram.forward(ai, nodeInputs, options);
       return {
         ...state,
         [`${nodeName}Result`]: result
@@ -15678,15 +15695,15 @@ var AxFlowTypedSubContextImpl = class {
   }
   steps = [];
   execute(nodeName, mapping, dynamicContext) {
-    const nodeGenerator = this.nodeGenerators.get(nodeName);
-    if (!nodeGenerator) {
-      throw new Error(`Node generator for '${nodeName}' not found.`);
+    const nodeProgram = this.nodeGenerators.get(nodeName);
+    if (!nodeProgram) {
+      throw new Error(`Node program for '${nodeName}' not found.`);
     }
     this.steps.push(async (state, context3) => {
       const ai = dynamicContext?.ai ?? context3.mainAi;
       const options = dynamicContext?.options ?? context3.mainOptions;
       const nodeInputs = mapping(state);
-      const result = await nodeGenerator.forward(ai, nodeInputs, options);
+      const result = await nodeProgram.forward(ai, nodeInputs, options);
       return {
         ...state,
         [`${nodeName}Result`]: result