@ax-llm/ax 11.0.52 → 11.0.54

package/index.cjs

@@ -66,6 +66,7 @@ __export(index_exports, {
   AxAIOpenAIResponsesImpl: () => AxAIOpenAIResponsesImpl,
   AxAIReka: () => AxAIReka,
   AxAIRekaModel: () => AxAIRekaModel,
+  AxAIServiceAbortedError: () => AxAIServiceAbortedError,
   AxAIServiceAuthenticationError: () => AxAIServiceAuthenticationError,
   AxAIServiceError: () => AxAIServiceError,
   AxAIServiceNetworkError: () => AxAIServiceNetworkError,
@@ -459,10 +460,25 @@ var AxAIServiceStreamTerminatedError = class extends AxAIServiceError {
 };
 var AxAIServiceTimeoutError = class extends AxAIServiceError {
   constructor(url, timeoutMs, requestBody, context3) {
-    super(`Request timeout after ${timeoutMs}ms`, url, requestBody, void 0, {
-      timeoutMs,
-      ...context3
-    });
+    super(
+      `Request timed out after ${timeoutMs}ms`,
+      url,
+      requestBody,
+      void 0,
+      { timeoutMs, ...context3 }
+    );
+    this.name = this.constructor.name;
+  }
+};
+var AxAIServiceAbortedError = class extends AxAIServiceError {
+  constructor(url, reason, requestBody, context3) {
+    super(
+      `Request aborted${reason ? `: ${reason}` : ""}`,
+      url,
+      requestBody,
+      void 0,
+      { abortReason: reason, ...context3 }
+    );
     this.name = this.constructor.name;
   }
 };
@@ -522,9 +538,34 @@ var apiCall = async (api, json) => {
   });
   let attempt = 0;
   while (true) {
-    const controller = new AbortController();
+    const combinedAbortController = new AbortController();
+    if (api.abortSignal) {
+      if (api.abortSignal.aborted) {
+        throw new AxAIServiceAbortedError(
+          apiUrl.href,
+          api.abortSignal.reason,
+          json,
+          { metrics }
+        );
+      }
+      const userAbortHandler = () => {
+        combinedAbortController.abort(
+          api.abortSignal.reason || "User aborted request"
+        );
+      };
+      api.abortSignal.addEventListener("abort", userAbortHandler, {
+        once: true
+      });
+      const originalAbort = combinedAbortController.abort.bind(
+        combinedAbortController
+      );
+      combinedAbortController.abort = (reason) => {
+        api.abortSignal.removeEventListener("abort", userAbortHandler);
+        originalAbort(reason);
+      };
+    }
     timeoutId = setTimeout(() => {
-      controller.abort("Request timeout");
+      combinedAbortController.abort("Request timeout");
     }, timeoutMs);
     try {
       const res = await (api.fetch ?? fetch)(apiUrl, {
@@ -536,7 +577,7 @@ var apiCall = async (api, json) => {
           ...api.headers
         },
         body: JSON.stringify(json),
-        signal: controller.signal
+        signal: combinedAbortController.signal
       });
       clearTimeout(timeoutId);
       if (res.status === 401 || res.status === 403) {
@@ -599,12 +640,12 @@ var apiCall = async (api, json) => {
       let lastChunk;
       let chunkCount = 0;
       const trackingStream = new import_web3.TransformStream({
-        transform(chunk, controller2) {
+        transform(chunk, controller) {
           lastChunk = chunk;
           chunkCount++;
           metrics.streamChunks = chunkCount;
           metrics.lastChunkTime = Date.now();
-          controller2.enqueue(chunk);
+          controller.enqueue(chunk);
           api.span?.addEvent("stream.chunk", {
             "stream.chunks": chunkCount,
             "stream.duration": Date.now() - metrics.startTime,
@@ -614,7 +655,7 @@ var apiCall = async (api, json) => {
       });
       let closed = false;
       return new import_web3.ReadableStream({
-        start(controller2) {
+        start(controller) {
           const reader = res.body.pipeThrough(new textDecoderStream()).pipeThrough(new SSEParser()).pipeThrough(trackingStream).getReader();
           async function read() {
             try {
@@ -623,12 +664,12 @@ var apiCall = async (api, json) => {
                 if (done) {
                   if (!closed) {
                     closed = true;
-                    controller2.close();
+                    controller.close();
                   }
                   break;
                 }
                 if (closed) break;
-                controller2.enqueue(value);
+                controller.enqueue(value);
               }
             } catch (e) {
               const error = e;
@@ -637,7 +678,7 @@ var apiCall = async (api, json) => {
                 streamDuration: Date.now() - metrics.startTime
               };
               if (error.name === "AbortError" || error.message?.includes("aborted")) {
-                controller2.error(
+                controller.error(
                   new AxAIServiceStreamTerminatedError(
                     apiUrl.href,
                     json,
@@ -646,7 +687,7 @@ var apiCall = async (api, json) => {
                   )
                 );
               } else if (error instanceof TypeError && error.message.includes("cancelled")) {
-                controller2.error(
+                controller.error(
                   new AxAIServiceStreamTerminatedError(
                     apiUrl.href,
                     json,
@@ -658,7 +699,7 @@ var apiCall = async (api, json) => {
                   )
                 );
               } else {
-                controller2.error(
+                controller.error(
                   new AxAIServiceNetworkError(
                     error,
                     apiUrl.href,
@@ -685,9 +726,18 @@ var apiCall = async (api, json) => {
       });
     } catch (error) {
       if (error instanceof Error && error.name === "AbortError") {
-        throw new AxAIServiceTimeoutError(apiUrl.href, timeoutMs, json, {
-          metrics
-        });
+        if (api.abortSignal?.aborted) {
+          throw new AxAIServiceAbortedError(
+            apiUrl.href,
+            api.abortSignal.reason,
+            json,
+            { metrics }
+          );
+        } else {
+          throw new AxAIServiceTimeoutError(apiUrl.href, timeoutMs, json, {
+            metrics
+          });
+        }
       }
       if (api.span?.isRecording()) {
         api.span.recordException(error);
@@ -932,6 +982,7 @@ var AxBaseAI = class {
   timeout;
   excludeContentFromTrace;
   models;
+  abortSignal;
   modelInfo;
   modelUsage;
   embedModelUsage;
@@ -992,6 +1043,7 @@ var AxBaseAI = class {
     this.timeout = options.timeout;
     this.tracer = options.tracer;
     this.excludeContentFromTrace = options.excludeContentFromTrace;
+    this.abortSignal = options.abortSignal;
   }
   getOptions() {
     return {
@@ -1000,7 +1052,8 @@ var AxBaseAI = class {
       fetch: this.fetch,
       tracer: this.tracer,
       timeout: this.timeout,
-      excludeContentFromTrace: this.excludeContentFromTrace
+      excludeContentFromTrace: this.excludeContentFromTrace,
+      abortSignal: this.abortSignal
     };
   }
   getModelList() {
@@ -1182,7 +1235,8 @@ var AxBaseAI = class {
           timeout: this.timeout,
           debug,
           fetch: this.fetch,
-          span
+          span,
+          abortSignal: options?.abortSignal ?? this.abortSignal
         },
         reqValue
       );
@@ -1339,7 +1393,8 @@ var AxBaseAI = class {
           debug,
           fetch: this.fetch,
           timeout: this.timeout,
-          span
+          span,
+          abortSignal: options?.abortSignal ?? this.abortSignal
         },
         reqValue
       );
@@ -3355,7 +3410,7 @@ var AxAIGoogleGeminiImpl = class {
       candidateCount: 1,
       stopSequences: req.modelConfig?.stopSequences ?? this.config.stopSequences,
       responseMimeType: "text/plain",
-      ...thinkingConfig ? { thinkingConfig } : {}
+      ...Object.keys(thinkingConfig).length > 0 ? { thinkingConfig } : {}
     };
     const safetySettings2 = this.config.safetySettings;
     const reqValue = {
@@ -5631,6 +5686,8 @@ var validateValue = (field, value) => {
   const ft = field.type ?? { name: "string", isArray: false };
   const validateSingleValue = (expectedType, val) => {
     switch (expectedType) {
+      case "class":
+        return typeof val === "string";
       case "code":
         return typeof val === "string";
       case "string":
@@ -5669,7 +5726,7 @@ var validateValue = (field, value) => {
     }
     if (msg) {
       throw new Error(
-        `Validation failed: Expected '${field.name}' to be a ${msg} instead got '${value}'`
+        `Validation failed: Expected '${field.name}' to be type '${msg}' instead got '${value}'`
       );
     }
     return;
@@ -5694,7 +5751,7 @@ var validateValue = (field, value) => {
     }
     if (msg) {
       throw new Error(
-        `Validation failed: Expected '${field.name}' to be a ${msg} instead got '${value}'`
+        `Validation failed: Expected '${field.name}' to be type '${msg}' instead got '${value}'`
       );
     }
     return;
@@ -5715,8 +5772,9 @@ var validateValue = (field, value) => {
     isValid = validateSingleValue(ft.name, value);
   }
   if (!isValid) {
+    const gotType = Array.isArray(value) ? "array" : typeof value;
     throw new Error(
-      `Validation failed: Expected '${field.name}' to be a ${field.type?.isArray ? "an array of " : ""}${ft.name} instead got '${typeof value}' (${value})`
+      `Validation failed: Expected '${field.name}' to be a ${field.type?.isArray ? "an array of " : ""}${ft.name} instead got '${gotType}' (${JSON.stringify(value)})`
     );
   }
 };
@@ -5921,11 +5979,15 @@ var AxPromptTemplate = class {
   task;
   thoughtFieldName;
   functions;
+  strictExamples;
+  optionalOutputFields;
   constructor(sig, options, fieldTemplates) {
     this.sig = sig;
     this.fieldTemplates = fieldTemplates;
     this.thoughtFieldName = options?.thoughtFieldName ?? "thought";
     this.functions = options?.functions;
+    this.strictExamples = options?.strictExamples ?? false;
+    this.optionalOutputFields = options?.optionalOutputFields ?? [];
     const task = [];
     const inArgs = renderDescFields(this.sig.getInputFields());
     const outArgs = renderDescFields(this.sig.getOutputFields());
@@ -6095,13 +6157,32 @@ ${outputFields}`);
   };
   renderExamples = (data) => {
     const list = [];
+    const inputExampleContext = {
+      isExample: true,
+      strictExamples: this.strictExamples,
+      optionalOutputFields: this.optionalOutputFields,
+      isInputField: true
+    };
+    const outputExampleContext = {
+      isExample: true,
+      strictExamples: this.strictExamples,
+      optionalOutputFields: this.optionalOutputFields,
+      isInputField: false
+    };
     for (const [index, item] of data.entries()) {
-      const renderedInputItem = this.sig.getInputFields().map((field) => this.renderInField(field, item)).filter((v) => v !== void 0).flat();
-      const renderedOutputItem = this.sig.getOutputFields().map((field) => this.renderInField(field, item)).filter((v) => v !== void 0).flat();
+      const renderedInputItem = this.sig.getInputFields().map((field) => this.renderInField(field, item, inputExampleContext)).filter((v) => v !== void 0).flat();
+      const outputFields = this.sig.getOutputFields();
+      const renderedOutputItem = outputFields.map((field) => this.renderInField(field, item, outputExampleContext)).filter((v) => v !== void 0).flat();
       if (renderedOutputItem.length === 0) {
-        throw new Error(
-          `Output fields are required in examples: index: ${index}, data: ${JSON.stringify(item)}`
+        const missingFields = outputFields.filter((field) => !item[field.name]);
+        const allMissingFieldsAreOptional = missingFields.every(
+          (field) => this.optionalOutputFields.includes(field.name)
         );
+        if (!allMissingFieldsAreOptional) {
+          throw new Error(
+            `Output fields are required in examples: index: ${index}, data: ${JSON.stringify(item)}`
+          );
+        }
       }
       const renderedItem = [...renderedInputItem, ...renderedOutputItem];
       if (index > 0 && renderedItem.length > 0 && renderedItem[0]?.type === "text") {
@@ -6121,9 +6202,26 @@ ${outputFields}`);
   };
   renderDemos = (data) => {
     const list = [];
-    const fields = [...this.sig.getInputFields(), ...this.sig.getOutputFields()];
+    const inputFields = this.sig.getInputFields();
+    const outputFields = this.sig.getOutputFields();
     for (const item of data) {
-      const renderedItem = fields.map((field) => this.renderInField(field, item)).filter((v) => v !== void 0).flat();
+      const inputRenderedItems = inputFields.map(
+        (field) => this.renderInField(field, item, {
+          isExample: true,
+          strictExamples: this.strictExamples,
+          optionalOutputFields: this.optionalOutputFields,
+          isInputField: true
+        })
+      ).filter((v) => v !== void 0).flat();
+      const outputRenderedItems = outputFields.map(
+        (field) => this.renderInField(field, item, {
+          isExample: true,
+          strictExamples: this.strictExamples,
+          optionalOutputFields: this.optionalOutputFields,
+          isInputField: false
+        })
+      ).filter((v) => v !== void 0).flat();
+      const renderedItem = [...inputRenderedItems, ...outputRenderedItems];
       renderedItem.slice(0, -1).forEach((v) => {
         if ("text" in v) {
           v.text = v.text + "\n";
@@ -6137,15 +6235,15 @@ ${outputFields}`);
     return list;
   };
   renderInputFields = (values) => {
-    const renderedItems = this.sig.getInputFields().map((field) => this.renderInField(field, values)).filter((v) => v !== void 0).flat();
+    const renderedItems = this.sig.getInputFields().map((field) => this.renderInField(field, values, void 0)).filter((v) => v !== void 0).flat();
     renderedItems.filter((v) => v.type === "text").forEach((v) => {
       v.text = v.text + "\n";
     });
     return renderedItems;
   };
-  renderInField = (field, values) => {
+  renderInField = (field, values, context3) => {
     const value = values[field.name];
-    if (isEmptyValue(field, value)) {
+    if (isEmptyValue(field, value, context3)) {
       return;
     }
     if (field.type) {
@@ -6329,15 +6427,34 @@ function combineConsecutiveStrings(separator) {
     return acc;
   };
 }
-var isEmptyValue = (field, value) => {
+var isEmptyValue = (field, value, context3) => {
   if (typeof value === "boolean") {
     return false;
   }
   if (!value || (Array.isArray(value) || typeof value === "string") && value.length === 0) {
-    if (field.isOptional || field.isInternal) {
+    if (context3?.isExample) {
+      const isInputField = context3?.isInputField ?? true;
+      if (isInputField) {
+        if (!context3?.strictExamples) {
+          return true;
+        } else {
+          if (field.isOptional || field.isInternal) {
+            return true;
+          }
+          throw new Error(`Value for input field '${field.name}' is required.`);
+        }
+      } else {
+        if (field.isOptional || field.isInternal || context3?.optionalOutputFields?.includes(field.name)) {
+          return true;
+        }
+        throw new Error(`Value for output field '${field.name}' is required.`);
+      }
+    }
+    if (field.isOptional || field.isInternal || context3?.optionalOutputFields?.includes(field.name)) {
       return true;
     }
-    throw new Error(`Value for input field '${field.name}' is required.`);
+    const fieldType = context3?.isInputField !== false ? "input" : "output";
+    throw new Error(`Value for ${fieldType} field '${field.name}' is required.`);
   }
   return false;
 };
@@ -7597,6 +7714,7 @@ var AxProgramWithSignature = class {
   signature;
   sigHash;
   examples;
+  examplesOptions;
   demos;
   trace;
   usage = [];
@@ -7638,16 +7756,16 @@ var AxProgramWithSignature = class {
       this.key.id = [parentId, this.key.id].join("/");
     }
   }
-  setExamples(examples) {
-    this._setExamples(examples);
+  setExamples(examples, options) {
+    this._setExamples(examples, options);
     if (!("programId" in examples)) {
       return;
     }
     for (const child of this.children) {
-      child.setExamples(examples);
+      child.setExamples(examples, options);
     }
   }
-  _setExamples(examples) {
+  _setExamples(examples, options) {
     let traces = [];
     if ("programId" in examples && examples.programId === this.key.id) {
       traces = examples.traces;
@@ -7656,6 +7774,7 @@ var AxProgramWithSignature = class {
       traces = examples;
     }
     if (traces) {
+      this.examplesOptions = options;
       const sig = this.signature;
       const fields = [...sig.getInputFields(), ...sig.getOutputFields()];
       this.examples = traces.map((e) => {
@@ -7736,12 +7855,12 @@ var AxProgram = class {
       this.key.id = [parentId, this.key.id].join("/");
     }
   }
-  setExamples(examples) {
+  setExamples(examples, options) {
     if (!("programId" in examples)) {
       return;
     }
     for (const child of this.children) {
-      child.setExamples(examples);
+      child.setExamples(examples, options);
     }
   }
   getTraces() {
@@ -7795,7 +7914,9 @@ var AxGen = class extends AxProgramWithSignature {
     this.thoughtFieldName = options?.thoughtFieldName ?? "thought";
     const promptTemplateOptions = {
       functions: options?.functions,
-      thoughtFieldName: this.thoughtFieldName
+      thoughtFieldName: this.thoughtFieldName,
+      strictExamples: options?.strictExamples,
+      optionalOutputFields: options?.optionalOutputFields
     };
     this.promptTemplate = new (options?.promptTemplate ?? AxPromptTemplate)(
       this.signature,
@@ -7881,7 +8002,8 @@ var AxGen = class extends AxProgramWithSignature {
         stream,
         debug: false,
         thinkingTokenBudget,
-        traceContext
+        traceContext,
+        abortSignal: options?.abortSignal
       }
     );
     return res;
@@ -8158,7 +8280,9 @@ Content: ${result.content}`
       const promptTemplateClass = this.options?.promptTemplate ?? AxPromptTemplate;
       const currentPromptTemplateOptions = {
         functions: options.functions,
-        thoughtFieldName: this.thoughtFieldName
+        thoughtFieldName: this.thoughtFieldName,
+        strictExamples: this.options?.strictExamples,
+        optionalOutputFields: this.options?.optionalOutputFields
       };
       this.promptTemplate = new promptTemplateClass(
         this.signature,
@@ -8344,6 +8468,22 @@ Content: ${result.content}`
       stream: true
     });
   }
+  setExamples(examples, options) {
+    super.setExamples(examples, options);
+    if (options?.optionalOutputFields) {
+      const promptTemplateClass = this.options?.promptTemplate ?? AxPromptTemplate;
+      const currentPromptTemplateOptions = {
+        functions: this.functions,
+        thoughtFieldName: this.thoughtFieldName,
+        strictExamples: this.options?.strictExamples,
+        optionalOutputFields: options.optionalOutputFields
+      };
+      this.promptTemplate = new promptTemplateClass(
+        this.signature,
+        currentPromptTemplateOptions
+      );
+    }
+  }
 };
 var AxGenerateError = class extends Error {
   details;
@@ -8474,8 +8614,8 @@ var AxAgent = class {
       this.func.parameters = addModelParameter(this.func.parameters, mm);
     }
   }
-  setExamples(examples) {
-    this.program.setExamples(examples);
+  setExamples(examples, options) {
+    this.program.setExamples(examples, options);
   }
   setId(id) {
     this.program.setId(id);
@@ -9769,7 +9909,12 @@ var AxDBManager = class {
       const bs = options?.batchSize ?? 10;
       for (let i = 0; i < chunks.length; i += bs) {
         const batch = chunks.slice(i, i + bs);
-        const ret = await this.ai.embed({ texts: batch });
+        const ret = await this.ai.embed(
+          { texts: batch },
+          {
+            abortSignal: options?.abortSignal
+          }
+        );
         const embeddings = ret.embeddings.map((embedding, index) => ({
           id: `chunk_${Date.now() + index}`,
           // Unique ID for each chunk, adjusted by index
@@ -9785,7 +9930,10 @@ var AxDBManager = class {
       throw new Error(`Error processing text: ${error}`);
     }
   };
-  query = async (query, { topPercent } = {}) => {
+  query = async (query, {
+    topPercent,
+    abortSignal
+  } = {}) => {
     const texts = Array.isArray(query) ? query : [query];
     if (typeof texts[0] === "string" && this.rewriter) {
       for (const [i, text] of texts.entries()) {
@@ -9797,7 +9945,12 @@ var AxDBManager = class {
     }
     let queries;
     if (typeof texts[0] === "string") {
-      const embedResults = await this.ai.embed({ texts });
+      const embedResults = await this.ai.embed(
+        { texts },
+        {
+          abortSignal
+        }
+      );
       queries = embedResults.embeddings.map(
         (values) => this.db.query({ table, values })
       );
@@ -11392,9 +11545,14 @@ var AxSimpleClassifier = class {
   setState(state) {
     this.db.setDB(state);
   }
-  setClasses = async (classes) => {
+  setClasses = async (classes, options) => {
     for (const c of classes) {
-      const ret = await this.ai.embed({ texts: c.getContext() });
+      const ret = await this.ai.embed(
+        { texts: c.getContext() },
+        {
+          abortSignal: options?.abortSignal
+        }
+      );
       await this.db.upsert({
         id: c.getName(),
         table: "classes",
@@ -11403,7 +11561,12 @@ var AxSimpleClassifier = class {
     }
   };
   async forward(text, options) {
-    const { embeddings } = await this.ai.embed({ texts: [text] });
+    const { embeddings } = await this.ai.embed(
+      { texts: [text] },
+      {
+        abortSignal: options?.abortSignal
+      }
+    );
     const matches = await this.db.query({
       table: "classes",
       values: embeddings[0]
@@ -11604,7 +11767,10 @@ var AxEmbeddingAdapter = class {
   async embedAdapter(text, extra) {
     const embedRes = await this.aiService.embed(
       { texts: [text] },
-      { sessionId: extra?.sessionId }
+      {
+        sessionId: extra?.sessionId,
+        abortSignal: extra?.abortSignal
+      }
     );
     const embeds = embedRes.embeddings.at(0);
     if (!embeds) {
@@ -13368,6 +13534,7 @@ var AxRAG = class extends AxChainOfThought {
   AxAIOpenAIResponsesImpl,
   AxAIReka,
   AxAIRekaModel,
+  AxAIServiceAbortedError,
   AxAIServiceAuthenticationError,
   AxAIServiceError,
   AxAIServiceNetworkError,