@ax-llm/ax 12.0.21 → 12.0.22

package/index.d.cts

@@ -3717,6 +3717,25 @@ TState extends AxFlowState = IN> extends AxProgramWithSignature<IN, OUT> {
         [K in TName]: InferAxGen<TSig>;
     }, // Add new node to registry
     TState>;
+    /**
+     * Declares a reusable computational node using an AxSignature instance.
+     * This allows using pre-configured signatures in the flow.
+     *
+     * @param name - The name of the node
+     * @param signature - AxSignature instance to use for this node
+     * @param options - Optional program forward options (same as AxGen)
+     * @returns New AxFlow instance with updated TNodes type
+     *
+     * @example
+     * ```typescript
+     * const sig = new AxSignature('text:string -> summary:string')
+     * flow.node('summarizer', sig, { temperature: 0.1 })
+     * ```
+     */
+    node<TName extends string>(name: TName, signature: AxSignature, options?: Readonly<AxProgramForwardOptions>): AxFlow<IN, OUT, TNodes & {
+        [K in TName]: AxGen<AxGenIn, AxGenOut>;
+    }, // Add new node to registry
+    TState>;
     /**
      * Declares a reusable computational node using an existing AxGen instance.
      * This allows reusing pre-configured generators in the flow.
@@ -3736,14 +3755,40 @@ TState extends AxFlowState = IN> extends AxProgramWithSignature<IN, OUT> {
     }, // Add new node to registry with exact type
     TState>;
     /**
-     * Short alias for node() - supports both signature strings and AxGen instances
+     * Declares a reusable computational node using a class that extends AxProgramWithSignature.
+     * This allows using custom program classes in the flow.
+     *
+     * @param name - The name of the node
+     * @param programClass - Class that extends AxProgramWithSignature to use for this node
+     * @returns New AxFlow instance with updated TNodes type
+     *
+     * @example
+     * ```typescript
+     * class CustomProgram extends AxProgramWithSignature<{ input: string }, { output: string }> {
+     *   async forward(ai, values) { return { output: values.input.toUpperCase() } }
+     * }
+     * flow.node('custom', CustomProgram)
+     * ```
+     */
+    node<TName extends string, TProgram extends new () => AxProgramWithSignature<any, any>>(name: TName, programClass: TProgram): AxFlow<IN, OUT, TNodes & {
+        [K in TName]: InstanceType<TProgram>;
+    }, // Add new node to registry with exact type
+    TState>;
+    /**
+     * Short alias for node() - supports signature strings, AxSignature instances, AxGen instances, and program classes
      */
     n<TName extends string, TSig extends string>(name: TName, signature: TSig, options?: Readonly<AxProgramForwardOptions>): AxFlow<IN, OUT, TNodes & {
         [K in TName]: InferAxGen<TSig>;
     }, TState>;
+    n<TName extends string>(name: TName, signature: AxSignature, options?: Readonly<AxProgramForwardOptions>): AxFlow<IN, OUT, TNodes & {
+        [K in TName]: AxGen<AxGenIn, AxGenOut>;
+    }, TState>;
     n<TName extends string, TGen extends AxGen<any, any>>(name: TName, axgenInstance: TGen): AxFlow<IN, OUT, TNodes & {
         [K in TName]: TGen;
     }, TState>;
+    n<TName extends string, TProgram extends new () => AxProgramWithSignature<any, any>>(name: TName, programClass: TProgram): AxFlow<IN, OUT, TNodes & {
+        [K in TName]: InstanceType<TProgram>;
+    }, TState>;
     /**
      * Applies a synchronous transformation to the state object.
      * Returns a new AxFlow type with the evolved state.
@@ -3951,7 +3996,7 @@ TState extends AxFlowState = IN> extends AxProgramWithSignature<IN, OUT> {
 declare class AxFlowTypedSubContextImpl<TNodes extends Record<string, AxGen<any, any>>, TState extends AxFlowState> implements AxFlowTypedSubContext<TNodes, TState> {
     private readonly nodeGenerators;
     private readonly steps;
-    constructor(nodeGenerators: Map<string, AxGen<AxGenIn, AxGenOut>>);
+    constructor(nodeGenerators: Map<string, AxGen<AxGenIn, AxGenOut> | AxProgramWithSignature<AxGenIn, AxGenOut>>);
     execute<TNodeName extends keyof TNodes & string>(nodeName: TNodeName, mapping: (state: TState) => GetGenIn<TNodes[TNodeName]>, dynamicContext?: AxFlowDynamicContext): AxFlowTypedSubContext<TNodes, AddNodeResult<TState, TNodeName, GetGenOut<TNodes[TNodeName]>>>;
     map<TNewState extends AxFlowState>(transform: (state: TState) => TNewState): AxFlowTypedSubContext<TNodes, TNewState>;
     executeSteps(initialState: TState, context: Readonly<{

package/index.d.ts

@@ -3717,6 +3717,25 @@ TState extends AxFlowState = IN> extends AxProgramWithSignature<IN, OUT> {
         [K in TName]: InferAxGen<TSig>;
     }, // Add new node to registry
     TState>;
+    /**
+     * Declares a reusable computational node using an AxSignature instance.
+     * This allows using pre-configured signatures in the flow.
+     *
+     * @param name - The name of the node
+     * @param signature - AxSignature instance to use for this node
+     * @param options - Optional program forward options (same as AxGen)
+     * @returns New AxFlow instance with updated TNodes type
+     *
+     * @example
+     * ```typescript
+     * const sig = new AxSignature('text:string -> summary:string')
+     * flow.node('summarizer', sig, { temperature: 0.1 })
+     * ```
+     */
+    node<TName extends string>(name: TName, signature: AxSignature, options?: Readonly<AxProgramForwardOptions>): AxFlow<IN, OUT, TNodes & {
+        [K in TName]: AxGen<AxGenIn, AxGenOut>;
+    }, // Add new node to registry
+    TState>;
     /**
      * Declares a reusable computational node using an existing AxGen instance.
      * This allows reusing pre-configured generators in the flow.
@@ -3736,14 +3755,40 @@ TState extends AxFlowState = IN> extends AxProgramWithSignature<IN, OUT> {
     }, // Add new node to registry with exact type
     TState>;
     /**
-     * Short alias for node() - supports both signature strings and AxGen instances
+     * Declares a reusable computational node using a class that extends AxProgramWithSignature.
+     * This allows using custom program classes in the flow.
+     *
+     * @param name - The name of the node
+     * @param programClass - Class that extends AxProgramWithSignature to use for this node
+     * @returns New AxFlow instance with updated TNodes type
+     *
+     * @example
+     * ```typescript
+     * class CustomProgram extends AxProgramWithSignature<{ input: string }, { output: string }> {
+     *   async forward(ai, values) { return { output: values.input.toUpperCase() } }
+     * }
+     * flow.node('custom', CustomProgram)
+     * ```
+     */
+    node<TName extends string, TProgram extends new () => AxProgramWithSignature<any, any>>(name: TName, programClass: TProgram): AxFlow<IN, OUT, TNodes & {
+        [K in TName]: InstanceType<TProgram>;
+    }, // Add new node to registry with exact type
+    TState>;
+    /**
+     * Short alias for node() - supports signature strings, AxSignature instances, AxGen instances, and program classes
      */
     n<TName extends string, TSig extends string>(name: TName, signature: TSig, options?: Readonly<AxProgramForwardOptions>): AxFlow<IN, OUT, TNodes & {
         [K in TName]: InferAxGen<TSig>;
     }, TState>;
+    n<TName extends string>(name: TName, signature: AxSignature, options?: Readonly<AxProgramForwardOptions>): AxFlow<IN, OUT, TNodes & {
+        [K in TName]: AxGen<AxGenIn, AxGenOut>;
+    }, TState>;
     n<TName extends string, TGen extends AxGen<any, any>>(name: TName, axgenInstance: TGen): AxFlow<IN, OUT, TNodes & {
         [K in TName]: TGen;
     }, TState>;
+    n<TName extends string, TProgram extends new () => AxProgramWithSignature<any, any>>(name: TName, programClass: TProgram): AxFlow<IN, OUT, TNodes & {
+        [K in TName]: InstanceType<TProgram>;
+    }, TState>;
     /**
      * Applies a synchronous transformation to the state object.
      * Returns a new AxFlow type with the evolved state.
@@ -3951,7 +3996,7 @@ TState extends AxFlowState = IN> extends AxProgramWithSignature<IN, OUT> {
 declare class AxFlowTypedSubContextImpl<TNodes extends Record<string, AxGen<any, any>>, TState extends AxFlowState> implements AxFlowTypedSubContext<TNodes, TState> {
     private readonly nodeGenerators;
     private readonly steps;
-    constructor(nodeGenerators: Map<string, AxGen<AxGenIn, AxGenOut>>);
+    constructor(nodeGenerators: Map<string, AxGen<AxGenIn, AxGenOut> | AxProgramWithSignature<AxGenIn, AxGenOut>>);
     execute<TNodeName extends keyof TNodes & string>(nodeName: TNodeName, mapping: (state: TState) => GetGenIn<TNodes[TNodeName]>, dynamicContext?: AxFlowDynamicContext): AxFlowTypedSubContext<TNodes, AddNodeResult<TState, TNodeName, GetGenOut<TNodes[TNodeName]>>>;
     map<TNewState extends AxFlowState>(transform: (state: TState) => TNewState): AxFlowTypedSubContext<TNodes, TNewState>;
     executeSteps(initialState: TState, context: Readonly<{

package/index.js

@@ -15008,18 +15008,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`
@@ -15030,11 +15043,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.
@@ -15117,15 +15134,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
@@ -15459,15 +15476,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
@@ -15493,15 +15510,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