@ax-llm/ax 10.0.34 → 10.0.36

package/index.js

@@ -405,6 +405,14 @@ var AxBaseAI = class {
       this.tracer = options.tracer;
     }
   }
+  getOptions() {
+    return {
+      debug: this.debug,
+      rateLimiter: this.rt,
+      fetch: this.fetch,
+      tracer: this.tracer
+    };
+  }
   getModelInfo() {
     const mi = getModelInfo({
       model: this.models.model,
@@ -702,11 +710,11 @@ var logChatRequest = (req) => {
         return `${colorLog.blueBright("System:")}
 ${colorLog.whiteBright(msg.content)}`;
       case "function":
-        return `${colorLog.blueBright("\nFunction Result:")}
+        return `${colorLog.blueBright("Function Result:")}
 ${colorLog.whiteBright(msg.result)}`;
       case "user": {
         if (typeof msg.content === "string") {
-          return `${colorLog.blueBright("\nUser:")}
+          return `${colorLog.blueBright("User:")}
 ${colorLog.whiteBright(msg.content)}`;
         }
         const items2 = msg.content.map((v) => {
@@ -719,7 +727,7 @@ ${colorLog.whiteBright(msg.content)}`;
               throw new Error("Invalid content type");
           }
         });
-        return `${colorLog.blueBright("\nUser:")}
+        return `${colorLog.blueBright("User:")}
 ${items2.join("\n")}`;
       }
       case "assistant": {
@@ -739,8 +747,7 @@ ${colorLog.whiteBright(msg.content ?? "<empty>")}`;
     }
   });
   if (items) {
-    console.log("\n==========");
-    console.log(items.join("\n"));
+    process.stdout.write("\n===\n" + items.join("\n") + "\n\n---\n");
   }
 };
 var logResponse = (resp) => {
@@ -2110,7 +2117,7 @@ var AxAIGoogleGeminiImpl = class {
       throw new Error("Chat prompt is empty");
     }
     const apiConfig = {
-      name: stream ? `/models/${model}:streamGenerateContent?alt=sse` : `/models/${model}:createContent`
+      name: stream ? `/models/${model}:streamGenerateContent?alt=sse` : `/models/${model}:generateContent`
     };
     if (this.isVertex === false) {
       const pf = stream ? "&" : "?";
@@ -3099,6 +3106,9 @@ var AxAI = class {
   setOptions(options) {
     this.ai.setOptions(options);
   }
+  getOptions() {
+    return this.ai.getOptions();
+  }
 };
 
 // prompts/agent.ts
@@ -3129,76 +3139,145 @@ function mergeFunctionCalls(functionCalls, functionCallDeltas) {
 }
 
 // mem/memory.ts
-var AxMemory = class {
-  data = [];
-  sdata = /* @__PURE__ */ new Map();
-  limit;
-  constructor(limit = 50) {
+var defaultLimit = 1e4;
+var MemoryImpl = class {
+  constructor(limit = defaultLimit) {
+    this.limit = limit;
     if (limit <= 0) {
       throw Error("argument 'limit' must be greater than 0");
     }
-    this.limit = limit;
   }
-  add(value, sessionId) {
-    const d = this.get(sessionId);
-    let n = 0;
+  data = [];
+  add(value) {
     if (Array.isArray(value)) {
-      n = d.push(...structuredClone(value));
+      this.data.push(...value.map((chat) => ({ chat: structuredClone(chat) })));
     } else {
-      n = d.push({
-        ...structuredClone(value)
+      this.data.push({
+        chat: structuredClone(value)
       });
     }
-    if (d.length > this.limit) {
-      d.splice(0, this.limit + n - this.limit);
+    if (this.data.length > this.limit) {
+      const removeCount = this.data.length - this.limit;
+      this.data.splice(0, removeCount);
     }
   }
-  addResult({ content, name, functionCalls }, sessionId) {
+  addResult({
+    content,
+    name,
+    functionCalls
+  }) {
     if (!content && (!functionCalls || functionCalls.length === 0)) {
       return;
     }
-    this.add({ content, name, role: "assistant", functionCalls }, sessionId);
+    this.add({ content, name, role: "assistant", functionCalls });
   }
-  updateResult({ content, name, functionCalls }, sessionId) {
-    const items = this.get(sessionId);
-    const lastItem = items.at(-1);
-    if (!lastItem || lastItem.role !== "assistant") {
-      this.addResult({ content, name, functionCalls }, sessionId);
+  updateResult({
+    content,
+    name,
+    functionCalls
+  }) {
+    const lastItem = this.data.at(-1);
+    if (!lastItem || lastItem.chat.role !== "assistant") {
+      this.addResult({ content, name, functionCalls });
       return;
     }
-    if ("content" in lastItem && content) {
-      lastItem.content = content;
+    if ("content" in lastItem.chat && content) {
+      lastItem.chat.content = content;
+    }
+    if ("name" in lastItem.chat && name) {
+      lastItem.chat.name = name;
     }
-    if ("name" in lastItem && name) {
-      lastItem.name = name;
+    if ("functionCalls" in lastItem.chat && functionCalls) {
+      lastItem.chat.functionCalls = functionCalls;
+    }
+  }
+  addTag(name) {
+    const lastItem = this.data.at(-1);
+    if (!lastItem) {
+      return;
+    }
+    if (!lastItem.tags) {
+      lastItem.tags = [];
+    }
+    if (!lastItem.tags.includes(name)) {
+      lastItem.tags.push(name);
+    }
+  }
+  rewindToTag(name) {
+    const tagIndex = this.data.findIndex((item) => item.tags?.includes(name));
+    if (tagIndex === -1) {
+      throw new Error(`Tag "${name}" not found`);
+    }
+    const removedItems = this.data.splice(tagIndex);
+    return removedItems.map((item) => item.chat);
+  }
+  removeByTag(name) {
+    const indices = this.data.reduce((acc, item, index) => {
+      if (item.tags?.includes(name)) {
+        acc.push(index);
+      }
+      return acc;
+    }, []);
+    if (indices.length === 0) {
+      throw new Error(`No items found with tag "${name}"`);
+    }
+    return indices.reverse().map((index) => this.data.splice(index, 1).at(0)?.chat).filter(Boolean).reverse();
+  }
+  history() {
+    return this.data.map((item) => item.chat);
+  }
+  getLast() {
+    const lastItem = this.data.at(-1);
+    return lastItem?.chat;
+  }
+  reset() {
+    this.data = [];
+  }
+};
+var AxMemory = class {
+  constructor(limit = defaultLimit) {
+    this.limit = limit;
+    this.defaultMemory = new MemoryImpl(limit);
+  }
+  memories = /* @__PURE__ */ new Map();
+  defaultMemory;
+  getMemory(sessionId) {
+    if (!sessionId) {
+      return this.defaultMemory;
     }
-    if ("functionCalls" in lastItem && functionCalls) {
-      lastItem.functionCalls = functionCalls;
+    if (!this.memories.has(sessionId)) {
+      this.memories.set(sessionId, new MemoryImpl(this.limit));
     }
+    return this.memories.get(sessionId);
+  }
+  add(value, sessionId) {
+    this.getMemory(sessionId).add(value);
+  }
+  addResult(result, sessionId) {
+    this.getMemory(sessionId).addResult(result);
+  }
+  updateResult(result, sessionId) {
+    this.getMemory(sessionId).updateResult(result);
+  }
+  addTag(name, sessionId) {
+    this.getMemory(sessionId).addTag(name);
+  }
+  rewindToTag(name, sessionId) {
+    return this.getMemory(sessionId).rewindToTag(name);
   }
   history(sessionId) {
-    return this.get(sessionId);
+    return this.getMemory(sessionId).history();
   }
   getLast(sessionId) {
-    const d = this.get(sessionId);
-    return d.at(-1);
+    return this.getMemory(sessionId).getLast();
   }
   reset(sessionId) {
     if (!sessionId) {
-      this.data = [];
+      this.defaultMemory.reset();
     } else {
-      this.sdata.set(sessionId, []);
+      this.memories.set(sessionId, new MemoryImpl(this.limit));
     }
   }
-  get(sessionId) {
-    if (!sessionId) {
-      return this.data;
-    }
-    if (!this.sdata.has(sessionId)) {
-      this.sdata.set(sessionId, []);
-    }
-    return this.sdata.get(sessionId) || [];
-  }
 };
 
 // dsp/asserts.ts
@@ -3218,12 +3297,11 @@ var AxAssertionError = class extends Error {
   }
   getValue = () => this.values;
   getOptional = () => this.optional;
-  // eslint-disable-next-line @typescript-eslint/no-unused-vars
-  getFixingInstructions = (_sig) => {
+  getFixingInstructions = () => {
     const extraFields = [];
     extraFields.push({
       name: "error",
-      title: "Error",
+      title: "Error In Output",
       description: this.message
     });
     return extraFields;
@@ -3232,17 +3310,12 @@ var AxAssertionError = class extends Error {
 var assertAssertions = (asserts, values) => {
   for (const assert of asserts) {
     const { fn, message, optional } = assert;
-    try {
-      const res = fn(values);
-      if (res === void 0) {
-        continue;
-      }
-      if (!res && message) {
-        throw new AxAssertionError({ message, values, optional });
-      }
-    } catch (e) {
-      const message2 = e.message;
-      throw new AxAssertionError({ message: message2, values, optional });
+    const res = fn(values);
+    if (res === void 0) {
+      continue;
+    }
+    if (!res && message) {
+      throw new AxAssertionError({ message, values, optional });
     }
   }
 };
@@ -3259,17 +3332,12 @@ var assertStreamingAssertions = (asserts, values, xstate, content, final) => {
   const currValue = content.substring(xstate.s);
   for (const assert of fieldAsserts) {
     const { message, optional, fn } = assert;
-    try {
-      const res = fn(currValue, final);
-      if (res === void 0) {
-        continue;
-      }
-      if (!res && message) {
-        throw new AxAssertionError({ message, values, optional });
-      }
-    } catch (e) {
-      const message2 = e.message;
-      throw new AxAssertionError({ message: message2, values, optional });
+    const res = fn(currValue, final);
+    if (res === void 0) {
+      continue;
+    }
+    if (!res && message) {
+      throw new AxAssertionError({ message, values, optional });
     }
   }
 };
@@ -4031,7 +4099,7 @@ var AxPromptTemplate = class {
     const inArgs = this.renderDescFields(this.sig.getInputFields());
     const outArgs = this.renderDescFields(this.sig.getOutputFields());
     const task = [
-      `You will be provided with the following fields: ${inArgs}. Your task is to generate two new fields: ${outArgs}.`
+      `You will be provided with the following fields: ${inArgs}. Your task is to generate new fields: ${outArgs}.`
     ];
     const funcs = functions?.map(
       (f) => "toFunction" in f ? f.toFunction() : f
@@ -4363,7 +4431,7 @@ function capitalizeFirstLetter(str) {
 }
 
 // dsp/validate.ts
-var AxValidationError = class extends Error {
+var ValidationError = class extends Error {
   field;
   value;
   constructor({
@@ -4383,9 +4451,9 @@ var AxValidationError = class extends Error {
     const f = this.field;
     const extraFields = [
       {
-        name: `invalidField`,
-        title: `Invalid Field`,
-        description: `Ensure the field \`${f.title}\` is of type \`${toFieldType(f.type)}\``
+        name: `outputError`,
+        title: `Invalid Output Field`,
+        description: `Invalid format for field \`${f.title}\` of type \`${toFieldType(f.type)}\`, format should match: \`${f.description}\``
       }
     ];
     return extraFields;
@@ -4398,7 +4466,7 @@ function parseLLMFriendlyDate(field, dateStr) {
     return _parseLLMFriendlyDate(dateStr);
   } catch (err) {
     const message = err.message;
-    throw new AxValidationError({ field, message, value: dateStr });
+    throw new ValidationError({ field, message, value: dateStr });
   }
 }
 function _parseLLMFriendlyDate(dateStr) {
@@ -4415,7 +4483,7 @@ function parseLLMFriendlyDateTime(field, dateStr) {
     return _parseLLMFriendlyDateTime(dateStr);
   } catch (err) {
     const message = err.message;
-    throw new AxValidationError({ field, message, value: dateStr });
+    throw new ValidationError({ field, message, value: dateStr });
   }
 }
 function _parseLLMFriendlyDateTime(dateTimeStr) {
@@ -4532,7 +4600,7 @@ var convertValueToType = (field, val) => {
 var expectedTypeError = (field, err, value = "") => {
   const exp = field.type?.isArray ? `array of ${field.type.name}` : field.type?.name;
   const message = `Error '${err.message}', expected '${exp}' got '${value}'`;
-  return new AxValidationError({ message, field, value });
+  return new ValidationError({ message, field, value });
 };
 function validateAndParseFieldValue(field, fieldValue) {
   const fv = fieldValue?.toLocaleLowerCase();
@@ -4748,6 +4816,7 @@ function parseFunctionCalls(ai, functionCalls, values, model) {
 }
 
 // dsp/generate.ts
+var colorLog4 = new ColorLog();
 var AxGen = class extends AxProgramWithSignature {
   promptTemplate;
   asserts;
@@ -4964,7 +5033,7 @@ var AxGen = class extends AxProgramWithSignature {
   }
   async _forward(ai, values, options, span) {
     const stopFunction = (options?.stopFunction ?? this.options?.stopFunction)?.toLowerCase();
-    const maxRetries = options?.maxRetries ?? this.options?.maxRetries ?? 15;
+    const maxRetries = options?.maxRetries ?? this.options?.maxRetries ?? 3;
     const maxSteps = options?.maxSteps ?? this.options?.maxSteps ?? 10;
     const mem = options?.mem ?? this.options?.mem ?? new AxMemory();
     let err;
@@ -4981,7 +5050,7 @@ var AxGen = class extends AxProgramWithSignature {
     });
     mem.add(prompt, options?.sessionId);
     multiStepLoop: for (let n = 0; n < maxSteps; n++) {
-      for (let i = 0; i < maxRetries; i++) {
+      for (let errCount = 0; errCount < maxRetries; errCount++) {
         try {
           const output = await this.forwardCore({
             options,
@@ -4989,35 +5058,49 @@ var AxGen = class extends AxProgramWithSignature {
             mem
           });
           const lastMemItem = mem.getLast(options?.sessionId);
-          const stopFunctionExecuted = stopFunction && this.functionsExecuted.has(stopFunction);
-          if (lastMemItem?.role === "function") {
-            if (!stopFunction || !stopFunctionExecuted) {
-              continue multiStepLoop;
+          if (lastMemItem) {
+            const stopFunctionExecuted = stopFunction && this.functionsExecuted.has(stopFunction);
+            if (lastMemItem.role === "function") {
+              if (!stopFunction || !stopFunctionExecuted) {
+                continue multiStepLoop;
+              }
+            }
+            if (!stopFunctionExecuted) {
+              assertRequiredFields(this.signature, output);
             }
-          }
-          if (!stopFunctionExecuted) {
-            assertRequiredFields(this.signature, output);
           }
           this.trace = { ...values, ...output };
           return output;
         } catch (e) {
-          let extraFields;
+          let errorFields;
           span?.recordException(e);
-          if (e instanceof AxValidationError) {
-            extraFields = e.getFixingInstructions();
+          if (e instanceof ValidationError) {
+            errorFields = e.getFixingInstructions();
             err = e;
           } else if (e instanceof AxAssertionError) {
             const e1 = e;
-            extraFields = e1.getFixingInstructions(this.signature);
+            errorFields = e1.getFixingInstructions();
             err = e;
           } else {
             throw e;
           }
-          if (extraFields) {
-            const content = this.promptTemplate.renderExtraFields(extraFields);
-            mem.add({ role: "user", content }, options?.sessionId);
-            if (options?.debug) {
-              console.log("Error Correction:", content);
+          if (errorFields) {
+            mem.add(
+              {
+                role: "user",
+                content: this.promptTemplate.renderExtraFields(errorFields)
+              },
+              options?.sessionId
+            );
+            mem.addTag("error");
+            if (ai.getOptions().debug) {
+              process.stdout.write(
+                colorLog4.red(
+                  `Error Correction:
+${JSON.stringify(errorFields, null, 2)}
+`
+                )
+              );
             }
           }
         }
@@ -5323,6 +5406,9 @@ var AxBalancer = class {
   setOptions(options) {
     this.currentService.setOptions(options);
   }
+  getOptions() {
+    return this.currentService.getOptions();
+  }
 };
 
 // dsp/optimize.ts
@@ -6528,7 +6614,7 @@ var AxJSInterpreter = class {
 };
 
 // dsp/router.ts
-var colorLog4 = new ColorLog();
+var colorLog5 = new ColorLog();
 var AxRoute = class {
   name;
   context;
@@ -6580,7 +6666,7 @@ var AxRouter = class {
     }
     if (this.debug) {
       console.log(
-        colorLog4.whiteBright(`query: ${text}`) + "\n" + colorLog4.greenBright(
+        colorLog5.whiteBright(`query: ${text}`) + "\n" + colorLog5.greenBright(
           JSON.stringify(m.map((m2) => `${m2.id}, ${m2.score}`))
         )
       );
@@ -6893,7 +6979,6 @@ export {
   AxRouter,
   AxSignature,
   AxSpanKindValues,
-  AxTestPrompt,
-  AxValidationError
+  AxTestPrompt
 };
 //# sourceMappingURL=index.js.map