@ax-llm/ax 12.0.12 → 12.0.13

package/index.cjs

@@ -6257,15 +6257,24 @@ var MemoryImpl = class {
       debugRequest(items, this.options?.debugHideSystemPrompt);
     }
   }
+  addFunctionResults(results) {
+    const chat = results.map(({ index, ...value }) => ({
+      index,
+      value: structuredClone(value)
+    }));
+    const lastItem = this.getLast();
+    if (lastItem?.role === "function") {
+      lastItem.chat.push(...chat);
+    } else {
+      this.data.push({ role: "function", chat });
+    }
+  }
   addResponse(results) {
-    const chat = results.map((result) => ({
-      index: result.index,
-      value: structuredClone(result)
+    const chat = results.map(({ index, ...value }) => ({
+      index,
+      value: structuredClone(value)
     }));
-    this.data.push({
-      role: "assistant",
-      chat
-    });
+    this.data.push({ role: "assistant", chat });
     if (this.options?.debug) {
       for (const result of results) {
         debugResponse(result);
@@ -6353,9 +6362,20 @@ var MemoryImpl = class {
   history(index) {
     const result = [];
     for (const { role, chat } of this.data) {
-      const value = chat.find((v) => v.index === index)?.value;
-      if (value) {
-        result.push({ role, ...value });
+      let values;
+      if (role === "function") {
+        values = chat.filter((v) => v.index === index).map((v) => v.value);
+      } else {
+        values = chat.find((v) => v.index === index)?.value;
+      }
+      if (Array.isArray(values)) {
+        result.push(
+          ...values.map(
+            (v) => ({ ...v, role })
+          )
+        );
+      } else if (values) {
+        result.push({ ...values, role });
       }
     }
     return result;
@@ -6393,20 +6413,8 @@ var AxMemory = class {
     axValidateChatResponseResult(results);
     this.getMemory(sessionId).addResponse(results);
   }
-  addFunctionResult({
-    functionId,
-    isError,
-    index,
-    result
-  }, sessionId) {
-    const functionMessage = {
-      role: "function",
-      functionId,
-      isError,
-      result
-    };
-    axValidateChatRequestMessage(functionMessage);
-    this.getMemory(sessionId).addRequest([functionMessage], index);
+  addFunctionResults(results, sessionId) {
+    this.getMemory(sessionId).addFunctionResults(results);
   }
   updateResult(result, sessionId) {
     this.getMemory(sessionId).updateResult(result);
@@ -7643,51 +7651,48 @@ var processFunctions = async ({
       }
       return {
         result: functionResult ?? "",
+        role: "function",
         functionId: func.id,
         index
       };
     }).catch((e) => {
-      if (e instanceof FunctionError) {
-        const result = e.getFixingInstructions();
-        if (span) {
-          const errorEventData = {
-            name: func.name,
-            message: e.toString()
-          };
-          if (!excludeContentFromTrace) {
-            errorEventData.args = func.args;
-            errorEventData.fixing_instructions = result;
-          }
-          span.addEvent("function.error", errorEventData);
+      if (!(e instanceof FunctionError)) {
+        throw e;
+      }
+      const result = e.getFixingInstructions();
+      if (span) {
+        const errorEventData = {
+          name: func.name,
+          message: e.toString()
+        };
+        if (!excludeContentFromTrace) {
+          errorEventData.args = func.args;
+          errorEventData.fixing_instructions = result;
         }
-        mem.addFunctionResult(
-          {
-            functionId: func.id,
-            isError: true,
-            index,
-            result
-          },
-          sessionId
-        );
-        mem.addTag("error", sessionId);
-        if (ai.getOptions().debug) {
-          const logger = ai.getLogger();
-          logger(`\u274C Function Error Correction:
+        span.addEvent("function.error", errorEventData);
+      }
+      if (ai.getOptions().debug) {
+        const logger = ai.getLogger();
+        logger(`\u274C Function Error Correction:
 ${result}`, {
-            tags: ["error"]
-          });
-        }
-      } else {
-        throw e;
+          tags: ["error"]
+        });
       }
+      return {
+        functionId: func.id,
+        isError: true,
+        index,
+        result,
+        role: "function"
+      };
     });
     return promise;
   });
   const results = await Promise.all(promises);
-  for (const result of results) {
-    if (result) {
-      mem.addFunctionResult(result, sessionId);
-    }
+  const functionResults = results.filter((result) => result !== void 0);
+  mem.addFunctionResults(functionResults, sessionId);
+  if (functionResults.some((result) => result.isError)) {
+    mem.addTag("error", sessionId);
   }
   return functionsExecuted;
 };
@@ -9961,6 +9966,96 @@ var AxProgram = class {
   }
 };
 
+// dsp/samples.ts
+function checkForFunctionCalls(mem, sessionId) {
+  const history = mem.history(0, sessionId);
+  const hasFunctionResults = history.some((msg) => msg.role === "function");
+  const hasFunctionCalls = history.some(
+    (msg) => msg.role === "assistant" && "functionCalls" in msg && Array.isArray(msg.functionCalls) && msg.functionCalls.length > 0
+  );
+  return hasFunctionCalls && hasFunctionResults;
+}
+function extractFunctionResults(mem, sessionId) {
+  const history = mem.history(0, sessionId);
+  const results = [];
+  const assistantMessages = history.filter(
+    (msg) => msg.role === "assistant" && "functionCalls" in msg && Array.isArray(msg.functionCalls) && msg.functionCalls.length > 0
+  );
+  const functionMessages = history.filter((msg) => msg.role === "function");
+  for (const assistantMsg of assistantMessages) {
+    if ("functionCalls" in assistantMsg && assistantMsg.functionCalls) {
+      for (const funcCall of assistantMsg.functionCalls) {
+        const funcResult = functionMessages.find(
+          (msg) => "functionId" in msg && msg.functionId === funcCall.id
+        );
+        if (funcResult && "result" in funcResult && "functionId" in funcResult) {
+          results.push({
+            index: results.length,
+            // Use sequential index for function results
+            functionName: funcCall.function.name,
+            functionId: funcCall.id,
+            args: funcCall.function.params || "",
+            result: String(funcResult.result),
+            isError: "isError" in funcResult ? Boolean(funcResult.isError) : false
+          });
+        }
+      }
+    }
+  }
+  return results;
+}
+async function selectFromSamples(buffer, options, mem, sessionId) {
+  if (!options?.resultPicker || buffer.length <= 1) {
+    return 0;
+  }
+  const resultPicker = options.resultPicker;
+  const hasFunctionCalls = mem ? checkForFunctionCalls(mem, sessionId) : false;
+  if (hasFunctionCalls && mem) {
+    const functionResults = extractFunctionResults(mem, sessionId);
+    const selectedIndex = await resultPicker({
+      type: "function",
+      results: functionResults
+    });
+    if (selectedIndex < 0 || selectedIndex >= functionResults.length) {
+      throw new Error(
+        `Result picker returned invalid index: ${selectedIndex}. Must be between 0 and ${functionResults.length - 1}`
+      );
+    }
+    return selectedIndex;
+  } else {
+    const fieldResults = buffer.map((b, index) => ({
+      index,
+      sample: b.delta
+    }));
+    const selectedIndex = await resultPicker({
+      type: "fields",
+      results: fieldResults
+    });
+    if (selectedIndex < 0 || selectedIndex >= buffer.length) {
+      throw new Error(
+        `Result picker returned invalid index: ${selectedIndex}. Must be between 0 and ${buffer.length - 1}`
+      );
+    }
+    return selectedIndex;
+  }
+}
+async function selectFromSamplesInMemory(mem, sessionId, options) {
+  const lastMemory = mem?.getLast(sessionId);
+  if (!lastMemory || lastMemory.role !== "assistant") {
+    return 0;
+  }
+  if (lastMemory.chat.length <= 1) {
+    return 0;
+  }
+  const buffer = lastMemory.chat.map((chat) => ({
+    version: 0,
+    index: chat.index,
+    delta: chat.value
+  }));
+  const selectedIndex = await selectFromSamples(buffer, options, mem, sessionId);
+  return selectedIndex;
+}
+
 // dsp/validate.ts
 function handleValidationError(mem, errorFields, ai, promptTemplate, sessionId) {
   mem.addRequest(
@@ -10076,7 +10171,10 @@ var AxGen = class extends AxProgramWithSignature {
       thinkingTokenBudget,
       showThoughts
     } = options ?? {};
-    const chatPrompt = mem?.history(0, sessionId) ?? [];
+    const selectedIndex = await selectFromSamplesInMemory(mem, sessionId, {
+      resultPicker: options?.resultPicker
+    });
+    const chatPrompt = mem?.history(selectedIndex, sessionId) ?? [];
     if (chatPrompt.length === 0) {
       throw new Error("No chat prompt found");
     }
@@ -10087,7 +10185,8 @@ var AxGen = class extends AxProgramWithSignature {
     }
     const modelConfig = {
       ...options?.modelConfig,
-      ...options?.sampleCount ? { n: options.sampleCount } : {}
+      ...options?.sampleCount ? { n: options.sampleCount } : {},
+      ...options?.sampleCount && options?.modelConfig?.temperature == 1 ? { temperature: 0.8 } : {}
     };
     const res = await ai.chat(
       {
@@ -10364,15 +10463,58 @@ var AxGen = class extends AxProgramWithSignature {
       currentVersion = delta.version;
       buffer = mergeDeltas(buffer, delta);
     }
-    const result = buffer[0]?.delta ?? {};
+    const selectedIndex = await selectFromSamples(
+      buffer,
+      {
+        resultPicker: options?.resultPicker
+      },
+      // Pass memory to enable function result selection
+      options?.mem,
+      options?.sessionId
+    );
+    const selectedResult = buffer[selectedIndex];
+    const result = selectedResult?.delta ?? {};
     this.trace = { ...values, ...result };
     return result;
   }
   async *streamingForward(ai, values, options) {
-    yield* this._forward1(ai, values, {
+    if (!options?.resultPicker) {
+      yield* this._forward1(ai, values, {
+        ...options,
+        stream: true
+      });
+      return;
+    }
+    const generator = this._forward1(ai, values, {
       ...options,
       stream: true
     });
+    let buffer = [];
+    let currentVersion = 0;
+    for await (const delta of generator) {
+      if (delta.version !== currentVersion) {
+        buffer = [];
+      }
+      currentVersion = delta.version;
+      buffer = mergeDeltas(buffer, delta);
+    }
+    const selectedIndex = await selectFromSamples(
+      buffer,
+      {
+        resultPicker: options?.resultPicker
+      },
+      // Pass memory to enable function result selection
+      options?.mem,
+      options?.sessionId
+    );
+    const selectedResult = buffer[selectedIndex];
+    if (selectedResult) {
+      yield {
+        version: currentVersion,
+        index: selectedIndex,
+        delta: selectedResult.delta
+      };
+    }
   }
   setExamples(examples, options) {
     super.setExamples(examples, options);