@ax-llm/ax 11.0.60 → 11.0.61

package/index.d.cts

@@ -2244,12 +2244,12 @@ type AxGenIn = {
     [key: string]: AxFieldValue;
 };
 type AxGenOut = Record<string, AxFieldValue>;
-type AxMessage = {
+type AxMessage<IN extends AxGenIn> = {
     role: 'user';
-    values: AxGenIn;
+    values: IN;
 } | {
     role: 'assistant';
-    values: AxGenIn;
+    values: IN;
 };
 
 type Writeable<T> = {
@@ -2271,7 +2271,7 @@ declare class AxPromptTemplate {
     private readonly thoughtFieldName;
     private readonly functions?;
     constructor(sig: Readonly<AxSignature>, options?: Readonly<AxPromptTemplateOptions>, fieldTemplates?: Record<string, AxFieldTemplateFn>);
-    render: <T extends AxGenIn>(values: T | ReadonlyArray<AxMessage>, // Allow T (AxGenIn) or array of AxMessages
+    render: <T extends AxGenIn>(values: T | ReadonlyArray<AxMessage<T>>, // Allow T (AxGenIn) or array of AxMessages
     { examples, demos, }: Readonly<{
         skipSystemPrompt?: boolean;
         examples?: Record<string, AxFieldValue>[];
@@ -2364,7 +2364,7 @@ type AxProgramUsage = AxChatResponse['modelUsage'] & {
 interface AxProgramWithSignatureOptions {
     description?: string;
 }
-declare class AxProgramWithSignature<IN extends AxGenIn | ReadonlyArray<AxMessage>, OUT extends AxGenOut> implements AxTunable, AxUsable {
+declare class AxProgramWithSignature<IN extends AxGenIn, OUT extends AxGenOut> implements AxTunable, AxUsable {
     protected signature: AxSignature;
     protected sigHash: string;
     protected examples?: Record<string, AxFieldValue>[];
@@ -2377,8 +2377,8 @@ declare class AxProgramWithSignature<IN extends AxGenIn | ReadonlyArray<AxMessag
     constructor(signature: Readonly<AxSignature | string>, options?: Readonly<AxProgramWithSignatureOptions>);
     getSignature(): AxSignature;
     register(prog: Readonly<AxTunable & AxUsable>): void;
-    forward(_ai: Readonly<AxAIService>, _values: IN, _options?: Readonly<AxProgramForwardOptions>): Promise<OUT>;
-    streamingForward(_ai: Readonly<AxAIService>, _values: IN, _options?: Readonly<AxProgramStreamingForwardOptions>): AxGenStreamingOut<OUT>;
+    forward(_ai: Readonly<AxAIService>, _values: IN | AxMessage<IN>[], _options?: Readonly<AxProgramForwardOptions>): Promise<OUT>;
+    streamingForward(_ai: Readonly<AxAIService>, _values: IN | AxMessage<IN>[], _options?: Readonly<AxProgramStreamingForwardOptions>): AxGenStreamingOut<OUT>;
     setId(id: string): void;
     setParentId(parentId: string): void;
     setExamples(examples: Readonly<AxProgramExamples>, options?: Readonly<AxSetExamplesOptions>): void;
@@ -2395,8 +2395,8 @@ declare class AxProgram<IN extends AxGenIn, OUT extends AxGenOut> implements AxT
     private children;
     constructor();
     register(prog: Readonly<AxTunable & AxUsable>): void;
-    forward(_ai: Readonly<AxAIService>, _values: IN, _options?: Readonly<AxProgramForwardOptions>): Promise<OUT>;
-    streamingForward(_ai: Readonly<AxAIService>, _values: IN, _options?: Readonly<AxProgramStreamingForwardOptions>): AxGenStreamingOut<OUT>;
+    forward(_ai: Readonly<AxAIService>, _values: IN | AxMessage<IN>[], _options?: Readonly<AxProgramForwardOptions>): Promise<OUT>;
+    streamingForward(_ai: Readonly<AxAIService>, _values: IN | AxMessage<IN>[], _options?: Readonly<AxProgramStreamingForwardOptions>): AxGenStreamingOut<OUT>;
     setId(id: string): void;
     setParentId(parentId: string): void;
     setExamples(examples: Readonly<AxProgramExamples>, options?: Readonly<AxSetExamplesOptions>): void;
@@ -2465,8 +2465,8 @@ declare class AxAgent<IN extends AxGenIn, OUT extends AxGenOut = AxGenOut> imple
      * Initializes the agent's execution context, processing child agents and their functions.
      */
     private init;
-    forward(parentAi: Readonly<AxAIService>, values: IN, options?: Readonly<AxProgramForwardOptions>): Promise<OUT>;
-    streamingForward(parentAi: Readonly<AxAIService>, values: IN, options?: Readonly<AxProgramStreamingForwardOptions>): AxGenStreamingOut<OUT>;
+    forward(parentAi: Readonly<AxAIService>, values: IN | AxMessage<IN>[], options?: Readonly<AxProgramForwardOptions>): Promise<OUT>;
+    streamingForward(parentAi: Readonly<AxAIService>, values: IN | AxMessage<IN>[], options?: Readonly<AxProgramStreamingForwardOptions>): AxGenStreamingOut<OUT>;
     /**
      * Updates the agent's description.
      * This updates both the stored description and the function's description.
@@ -2840,7 +2840,7 @@ interface AxStreamingEvent<T> {
         functions?: AxChatResponseFunctionCall[];
     };
 }
-declare class AxGen<IN extends AxGenIn | ReadonlyArray<AxMessage> = AxGenIn | ReadonlyArray<AxMessage>, OUT extends AxGenerateResult<AxGenOut> = AxGenerateResult<AxGenOut>> extends AxProgramWithSignature<IN, OUT> {
+declare class AxGen<IN extends AxGenIn, OUT extends AxGenerateResult<AxGenOut> = AxGenerateResult<AxGenOut>> extends AxProgramWithSignature<IN, OUT> {
     private promptTemplate;
     private asserts;
     private streamingAsserts;
@@ -2865,12 +2865,12 @@ declare class AxGen<IN extends AxGenIn | ReadonlyArray<AxMessage> = AxGenIn | Re
     private processResponse;
     private _forward2;
     private shouldContinueSteps;
-    _forward1(ai: Readonly<AxAIService>, values: IN, options: Readonly<AxProgramForwardOptions>): AsyncGenerator<{
+    _forward1(ai: Readonly<AxAIService>, values: IN | AxMessage<IN>[], options: Readonly<AxProgramForwardOptions>): AsyncGenerator<{
         version: number;
         delta: Partial<OUT>;
     }, void, unknown>;
-    forward(ai: Readonly<AxAIService>, values: IN, options?: Readonly<AxProgramForwardOptions>): Promise<OUT>;
-    streamingForward(ai: Readonly<AxAIService>, values: IN, options?: Readonly<AxProgramStreamingForwardOptions>): AsyncGenerator<{
+    forward(ai: Readonly<AxAIService>, values: IN | AxMessage<IN>[], options?: Readonly<AxProgramForwardOptions>): Promise<OUT>;
+    streamingForward(ai: Readonly<AxAIService>, values: IN | AxMessage<IN>[], options?: Readonly<AxProgramStreamingForwardOptions>): AsyncGenerator<{
         version: number;
         delta: Partial<OUT>;
     }, void, unknown>;
@@ -3690,9 +3690,13 @@ declare class AxRAG extends AxChainOfThought<{
     constructor(queryFn: (query: string) => Promise<string>, options: Readonly<AxProgramForwardOptions & {
         maxHops?: number;
     }>);
-    forward(ai: Readonly<AxAIService>, { question }: Readonly<{
+    forward(ai: Readonly<AxAIService>, values: {
+        context: string[];
         question: string;
-    }>, options?: Readonly<AxProgramForwardOptions>): Promise<{
+    } | AxMessage<{
+        context: string[];
+        question: string;
+    }>[], options?: Readonly<AxProgramForwardOptions>): Promise<{
         answer: string;
     }>;
 }

package/index.d.ts

@@ -2244,12 +2244,12 @@ type AxGenIn = {
     [key: string]: AxFieldValue;
 };
 type AxGenOut = Record<string, AxFieldValue>;
-type AxMessage = {
+type AxMessage<IN extends AxGenIn> = {
     role: 'user';
-    values: AxGenIn;
+    values: IN;
 } | {
     role: 'assistant';
-    values: AxGenIn;
+    values: IN;
 };
 
 type Writeable<T> = {
@@ -2271,7 +2271,7 @@ declare class AxPromptTemplate {
     private readonly thoughtFieldName;
     private readonly functions?;
     constructor(sig: Readonly<AxSignature>, options?: Readonly<AxPromptTemplateOptions>, fieldTemplates?: Record<string, AxFieldTemplateFn>);
-    render: <T extends AxGenIn>(values: T | ReadonlyArray<AxMessage>, // Allow T (AxGenIn) or array of AxMessages
+    render: <T extends AxGenIn>(values: T | ReadonlyArray<AxMessage<T>>, // Allow T (AxGenIn) or array of AxMessages
     { examples, demos, }: Readonly<{
         skipSystemPrompt?: boolean;
         examples?: Record<string, AxFieldValue>[];
@@ -2364,7 +2364,7 @@ type AxProgramUsage = AxChatResponse['modelUsage'] & {
 interface AxProgramWithSignatureOptions {
     description?: string;
 }
-declare class AxProgramWithSignature<IN extends AxGenIn | ReadonlyArray<AxMessage>, OUT extends AxGenOut> implements AxTunable, AxUsable {
+declare class AxProgramWithSignature<IN extends AxGenIn, OUT extends AxGenOut> implements AxTunable, AxUsable {
     protected signature: AxSignature;
     protected sigHash: string;
     protected examples?: Record<string, AxFieldValue>[];
@@ -2377,8 +2377,8 @@ declare class AxProgramWithSignature<IN extends AxGenIn | ReadonlyArray<AxMessag
     constructor(signature: Readonly<AxSignature | string>, options?: Readonly<AxProgramWithSignatureOptions>);
     getSignature(): AxSignature;
     register(prog: Readonly<AxTunable & AxUsable>): void;
-    forward(_ai: Readonly<AxAIService>, _values: IN, _options?: Readonly<AxProgramForwardOptions>): Promise<OUT>;
-    streamingForward(_ai: Readonly<AxAIService>, _values: IN, _options?: Readonly<AxProgramStreamingForwardOptions>): AxGenStreamingOut<OUT>;
+    forward(_ai: Readonly<AxAIService>, _values: IN | AxMessage<IN>[], _options?: Readonly<AxProgramForwardOptions>): Promise<OUT>;
+    streamingForward(_ai: Readonly<AxAIService>, _values: IN | AxMessage<IN>[], _options?: Readonly<AxProgramStreamingForwardOptions>): AxGenStreamingOut<OUT>;
     setId(id: string): void;
     setParentId(parentId: string): void;
     setExamples(examples: Readonly<AxProgramExamples>, options?: Readonly<AxSetExamplesOptions>): void;
@@ -2395,8 +2395,8 @@ declare class AxProgram<IN extends AxGenIn, OUT extends AxGenOut> implements AxT
     private children;
     constructor();
     register(prog: Readonly<AxTunable & AxUsable>): void;
-    forward(_ai: Readonly<AxAIService>, _values: IN, _options?: Readonly<AxProgramForwardOptions>): Promise<OUT>;
-    streamingForward(_ai: Readonly<AxAIService>, _values: IN, _options?: Readonly<AxProgramStreamingForwardOptions>): AxGenStreamingOut<OUT>;
+    forward(_ai: Readonly<AxAIService>, _values: IN | AxMessage<IN>[], _options?: Readonly<AxProgramForwardOptions>): Promise<OUT>;
+    streamingForward(_ai: Readonly<AxAIService>, _values: IN | AxMessage<IN>[], _options?: Readonly<AxProgramStreamingForwardOptions>): AxGenStreamingOut<OUT>;
     setId(id: string): void;
     setParentId(parentId: string): void;
     setExamples(examples: Readonly<AxProgramExamples>, options?: Readonly<AxSetExamplesOptions>): void;
@@ -2465,8 +2465,8 @@ declare class AxAgent<IN extends AxGenIn, OUT extends AxGenOut = AxGenOut> imple
      * Initializes the agent's execution context, processing child agents and their functions.
      */
     private init;
-    forward(parentAi: Readonly<AxAIService>, values: IN, options?: Readonly<AxProgramForwardOptions>): Promise<OUT>;
-    streamingForward(parentAi: Readonly<AxAIService>, values: IN, options?: Readonly<AxProgramStreamingForwardOptions>): AxGenStreamingOut<OUT>;
+    forward(parentAi: Readonly<AxAIService>, values: IN | AxMessage<IN>[], options?: Readonly<AxProgramForwardOptions>): Promise<OUT>;
+    streamingForward(parentAi: Readonly<AxAIService>, values: IN | AxMessage<IN>[], options?: Readonly<AxProgramStreamingForwardOptions>): AxGenStreamingOut<OUT>;
     /**
      * Updates the agent's description.
      * This updates both the stored description and the function's description.
@@ -2840,7 +2840,7 @@ interface AxStreamingEvent<T> {
         functions?: AxChatResponseFunctionCall[];
     };
 }
-declare class AxGen<IN extends AxGenIn | ReadonlyArray<AxMessage> = AxGenIn | ReadonlyArray<AxMessage>, OUT extends AxGenerateResult<AxGenOut> = AxGenerateResult<AxGenOut>> extends AxProgramWithSignature<IN, OUT> {
+declare class AxGen<IN extends AxGenIn, OUT extends AxGenerateResult<AxGenOut> = AxGenerateResult<AxGenOut>> extends AxProgramWithSignature<IN, OUT> {
     private promptTemplate;
     private asserts;
     private streamingAsserts;
@@ -2865,12 +2865,12 @@ declare class AxGen<IN extends AxGenIn | ReadonlyArray<AxMessage> = AxGenIn | Re
     private processResponse;
     private _forward2;
     private shouldContinueSteps;
-    _forward1(ai: Readonly<AxAIService>, values: IN, options: Readonly<AxProgramForwardOptions>): AsyncGenerator<{
+    _forward1(ai: Readonly<AxAIService>, values: IN | AxMessage<IN>[], options: Readonly<AxProgramForwardOptions>): AsyncGenerator<{
         version: number;
         delta: Partial<OUT>;
     }, void, unknown>;
-    forward(ai: Readonly<AxAIService>, values: IN, options?: Readonly<AxProgramForwardOptions>): Promise<OUT>;
-    streamingForward(ai: Readonly<AxAIService>, values: IN, options?: Readonly<AxProgramStreamingForwardOptions>): AsyncGenerator<{
+    forward(ai: Readonly<AxAIService>, values: IN | AxMessage<IN>[], options?: Readonly<AxProgramForwardOptions>): Promise<OUT>;
+    streamingForward(ai: Readonly<AxAIService>, values: IN | AxMessage<IN>[], options?: Readonly<AxProgramStreamingForwardOptions>): AsyncGenerator<{
         version: number;
         delta: Partial<OUT>;
     }, void, unknown>;
@@ -3690,9 +3690,13 @@ declare class AxRAG extends AxChainOfThought<{
     constructor(queryFn: (query: string) => Promise<string>, options: Readonly<AxProgramForwardOptions & {
         maxHops?: number;
     }>);
-    forward(ai: Readonly<AxAIService>, { question }: Readonly<{
+    forward(ai: Readonly<AxAIService>, values: {
+        context: string[];
         question: string;
-    }>, options?: Readonly<AxProgramForwardOptions>): Promise<{
+    } | AxMessage<{
+        context: string[];
+        question: string;
+    }>[], options?: Readonly<AxProgramForwardOptions>): Promise<{
         answer: string;
     }>;
 }

package/index.js

@@ -8541,9 +8541,21 @@ function processChildAgentFunction(childFunction, parentValues, parentInputKeys,
       );
       const originalFunc = processedFunction.func;
       processedFunction.func = async (childArgs, funcOptions) => {
+        let valuesToInject = {};
+        if (Array.isArray(parentValues)) {
+          const lastUserMessage = parentValues.filter((msg) => msg.role === "user").pop();
+          if (lastUserMessage) {
+            valuesToInject = pick(
+              lastUserMessage.values,
+              injectionKeys
+            );
+          }
+        } else {
+          valuesToInject = pick(parentValues, injectionKeys);
+        }
         const updatedChildArgs = {
           ...childArgs,
-          ...pick(parentValues, injectionKeys)
+          ...valuesToInject
         };
         if (options.debug && injectionKeys.length > 0) {
           const ai = funcOptions?.ai;
@@ -13532,7 +13544,17 @@ var AxRAG = class extends AxChainOfThought {
     this.queryFn = queryFn;
     this.register(this.genQuery);
   }
-  async forward(ai, { question }, options) {
+  async forward(ai, values, options) {
+    let question;
+    if (Array.isArray(values)) {
+      const lastUserMessage = values.filter((msg) => msg.role === "user").pop();
+      if (!lastUserMessage) {
+        throw new Error("No user message found in values array");
+      }
+      question = lastUserMessage.values.question;
+    } else {
+      question = values.question;
+    }
     let hop = 0;
     let context3 = [];
     while (hop < this.maxHops) {