@ax-llm/ax 10.0.39 → 10.0.40

package/index.js

@@ -525,9 +525,11 @@ var AxBaseAI = class {
           }
         },
         async (span) => {
-          const res = await this._chat2(model, modelConfig, req, options, span);
-          span.end();
-          return res;
+          try {
+            return await this._chat2(model, modelConfig, req, options, span);
+          } finally {
+            span.end();
+          }
         }
       );
     }
@@ -650,9 +652,11 @@ var AxBaseAI = class {
           }
         },
         async (span) => {
-          const res = await this._embed2(embedModel, req, options, span);
-          span.end();
-          return res;
+          try {
+            return await this._embed2(embedModel, req, options, span);
+          } finally {
+            span.end();
+          }
         }
       );
     }
@@ -720,7 +724,7 @@ ${colorLog.whiteBright(msg.content)}`;
         const items2 = msg.content.map((v) => {
           switch (v.type) {
             case "text":
-              return `(Text) ${colorLog.whiteBright(v.text)}`;
+              return `${colorLog.whiteBright(v.text)}`;
             case "image":
               return `(Image, ${v.mimeType}) ${colorLog.whiteBright(v.image.substring(0, 10))}`;
             default:
@@ -875,6 +879,30 @@ var AxAIAnthropicImpl = class {
     const apiConfig = {
       name: "/messages"
     };
+    let toolsChoice;
+    if (req.functionCall && req.functions && req.functions.length > 0) {
+      if (typeof req.functionCall === "string") {
+        switch (req.functionCall) {
+          case "auto":
+            toolsChoice = { tool_choice: { type: "auto" } };
+            break;
+          case "required":
+            toolsChoice = { tool_choice: { type: "any" } };
+            break;
+          case "none":
+            throw new Error("functionCall none not supported");
+        }
+      } else if ("function" in req.functionCall) {
+        toolsChoice = {
+          tool_choice: {
+            type: "tool",
+            name: req.functionCall.function.name
+          }
+        };
+      } else {
+        throw new Error("Invalid function call type, must be string or object");
+      }
+    }
     const system = req.chatPrompt.filter((msg) => msg.role === "system").map((msg) => ({
       type: "text",
       text: msg.content,
@@ -897,7 +925,8 @@ var AxAIAnthropicImpl = class {
       temperature: req.modelConfig?.temperature ?? this.config.temperature,
       top_p: req.modelConfig?.topP ?? this.config.topP,
       top_k: req.modelConfig?.topK ?? this.config.topK,
-      ...tools && tools.length > 0 ? { tools, tool_choice: { type: "auto" } } : {},
+      ...toolsChoice,
+      ...tools && tools.length > 0 ? { tools } : {},
       ...stream ? { stream: true } : {},
       ...system ? { system } : {},
       messages
@@ -3206,9 +3235,6 @@ var AxAI = class {
   }
 };
 
-// prompts/agent.ts
-import { SpanKind as SpanKind3 } from "@opentelemetry/api";
-
 // dsp/generate.ts
 import { ReadableStream } from "stream/web";
 import { SpanKind as SpanKind2 } from "@opentelemetry/api";
@@ -3397,7 +3423,7 @@ var AxAssertionError = class extends Error {
     extraFields.push({
       name: "error",
       title: "Error In Output",
-      description: this.message
+      description: `You must follow the following instructions, "${this.message}".`
     });
     return extraFields;
   };
@@ -3436,18 +3462,6 @@ var assertStreamingAssertions = (asserts, values, xstate, content, final) => {
     }
   }
 };
-var assertRequiredFields = (sig, values) => {
-  const fields = sig.getOutputFields();
-  const missingFields = fields.filter(
-    (f) => !f.isOptional && !(f.name in values)
-  );
-  if (missingFields.length > 0) {
-    throw new AxAssertionError({
-      message: `Output must include: ${missingFields.map((f) => `\`${f.title}:\``).join(", ")}`,
-      values
-    });
-  }
-};
 
 // dsp/extract.ts
 import JSON5 from "json5";
@@ -3995,6 +4009,21 @@ var parseMarkdownList = (input) => {
   }
   return list;
 };
+function mergeDeltas(base, delta) {
+  const merged = { ...base };
+  for (const key in delta) {
+    const baseValue = base[key];
+    const deltaValue = delta[key];
+    if (Array.isArray(baseValue) && Array.isArray(deltaValue)) {
+      merged[key] = [...baseValue, ...deltaValue];
+    } else if (typeof baseValue === "string" && typeof deltaValue === "string") {
+      merged[key] = baseValue + deltaValue;
+    } else {
+      merged[key] = deltaValue;
+    }
+  }
+  return merged;
+}
 
 // dsp/program.ts
 var AxProgramWithSignature = class {
@@ -4027,6 +4056,9 @@ var AxProgramWithSignature = class {
   async forward(_ai, _values, _options) {
     throw new Error("forward() not implemented");
   }
+  async *streamingForward(_ai, _values, _options) {
+    throw new Error("streamingForward() not implemented");
+  }
   setId(id) {
     this.key = { id, custom: true };
     for (const child of this.children) {
@@ -4121,6 +4153,9 @@ var AxProgram = class {
   async forward(_ai, _values, _options) {
     throw new Error("forward() not implemented");
   }
+  async *streamingForward(_ai, _values, _options) {
+    throw new Error("streamingForward() not implemented");
+  }
   setId(id) {
     this.key = { id, custom: true };
     for (const child of this.children) {
@@ -4218,9 +4253,10 @@ ${outputFields}`);
     task.push(formattingRules.trim());
     const desc = this.sig.getDescription();
     if (desc) {
+      const capitalized = capitalizeFirstLetter(desc.trim());
       task.push(
         `## TASK DESCRIPTION
-${capitalizeFirstLetter(desc.endsWith(".") ? desc : desc + ".")}`
+${capitalized.endsWith(".") ? capitalized : capitalized + "."}`
       );
     }
     this.task = {
@@ -4268,16 +4304,42 @@ ${capitalizeFirstLetter(desc.endsWith(".") ? desc : desc + ".")}`
   };
   renderExtraFields = (extraFields) => {
     const prompt = [];
-    if (extraFields && extraFields.length > 0) {
-      extraFields.forEach((field) => {
-        const fn = this.fieldTemplates?.[field.name] ?? this.defaultRenderInField;
-        prompt.push(...fn(field, field.description));
-      });
-    }
-    if (prompt.every((v) => v.type === "text")) {
-      return prompt.map((v) => v.text).join("\n\n");
-    }
-    return prompt.reduce(combineConsecutiveStrings("\n"), []);
+    if (!extraFields || extraFields.length === 0) {
+      return prompt;
+    }
+    const groupedFields = extraFields.reduce(
+      (acc, field) => {
+        const title = field.title;
+        if (!acc[title]) {
+          acc[title] = [];
+        }
+        acc[title].push(field);
+        return acc;
+      },
+      {}
+    );
+    const formattedGroupedFields = Object.entries(groupedFields).map(([title, fields]) => {
+      if (fields.length === 1) {
+        const field = fields[0];
+        return {
+          title,
+          name: field.name,
+          description: field.description
+        };
+      } else if (fields.length > 1) {
+        const valuesList = fields.map((field) => `- ${field.description}`).join("\n");
+        return {
+          title,
+          name: fields[0].name,
+          description: valuesList
+        };
+      }
+    }).filter(Boolean);
+    formattedGroupedFields.forEach((field) => {
+      const fn = this.fieldTemplates?.[field.name] ?? this.defaultRenderInField;
+      prompt.push(...fn(field, field.description));
+    });
+    return prompt;
   };
   renderExamples = (data) => {
     const list = [];
@@ -4526,34 +4588,47 @@ function capitalizeFirstLetter(str) {
 }
 
 // dsp/validate.ts
+var colorLog4 = new ColorLog();
 var ValidationError = class extends Error {
-  field;
-  value;
+  fields;
   constructor({
     message,
-    field,
-    value
+    fields
   }) {
     super(message);
-    this.field = field;
-    this.value = value;
+    this.fields = fields;
     this.name = this.constructor.name;
     Error.captureStackTrace(this, this.constructor);
   }
-  getField = () => this.field;
-  getValue = () => this.value;
+  getFields = () => this.fields;
   getFixingInstructions = () => {
-    const f = this.field;
-    const extraFields = [
-      {
-        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;
+    return this.fields.map((field) => ({
+      name: "outputError",
+      title: "Error In Output",
+      description: `Please fix and return the field \`${field.title}\` of type \`${toFieldType(field.type)}\`, ${this.message}.`
+    }));
   };
 };
+function handleValidationError(mem, errorFields, ai, promptTemplate, sessionId) {
+  mem.add(
+    {
+      role: "user",
+      content: promptTemplate.renderExtraFields(errorFields)
+    },
+    sessionId
+  );
+  mem.addTag("error");
+  if (ai.getOptions().debug) {
+    process.stdout.write(
+      colorLog4.red(
+        `
+Error Correction:
+${JSON.stringify(errorFields, null, 2)}
+`
+      )
+    );
+  }
+}
 
 // dsp/datetime.ts
 function parseLLMFriendlyDate(field, dateStr) {
@@ -4561,7 +4636,7 @@ function parseLLMFriendlyDate(field, dateStr) {
     return _parseLLMFriendlyDate(dateStr);
   } catch (err) {
     const message = err.message;
-    throw new ValidationError({ field, message, value: dateStr });
+    throw new ValidationError({ fields: [field], message, value: dateStr });
   }
 }
 function _parseLLMFriendlyDate(dateStr) {
@@ -4578,7 +4653,7 @@ function parseLLMFriendlyDateTime(field, dateStr) {
     return _parseLLMFriendlyDateTime(dateStr);
   } catch (err) {
     const message = err.message;
-    throw new ValidationError({ field, message, value: dateStr });
+    throw new ValidationError({ fields: [field], message, value: dateStr });
   }
 }
 function _parseLLMFriendlyDateTime(dateTimeStr) {
@@ -4620,41 +4695,61 @@ var formatDateWithTimezone = (date) => {
 
 // dsp/extract.ts
 var extractValues = (sig, values, content) => {
-  const xstate = { s: -1 };
+  const xstate = { extractedFields: [], s: -1 };
   streamingExtractValues(sig, values, xstate, content);
-  streamingExtractFinalValue(values, xstate, content);
+  streamingExtractFinalValue(sig, values, xstate, content);
+};
+var checkMissingRequiredFields = (xstate, values, currentIndex) => {
+  const missingFields = [];
+  for (let i = 0; i < currentIndex; i++) {
+    const field = xstate.extractedFields[i];
+    if (field && !field.isOptional && values[field.name] === void 0) {
+      missingFields.push(field);
+    }
+  }
+  if (missingFields.length > 0) {
+    throw new ValidationError({
+      message: `Required ${missingFields.length === 1 ? "field" : "fields"} not found`,
+      fields: missingFields
+    });
+  }
 };
-var streamingExtractValues = (sig, values, state, content) => {
+var streamingExtractValues = (sig, values, xstate, content) => {
   const fields = sig.getOutputFields();
-  for (const field of fields) {
+  for (const [index, field] of fields.entries()) {
     if (field.name in values) {
       continue;
     }
     const prefix = field.title + ":";
-    const e = content.indexOf(prefix, state.s + 1);
+    const e = content.indexOf(prefix, xstate.s + 1);
     if (e === -1) {
       continue;
     }
-    if (state.currField) {
-      const val = content.substring(state.s, e).trim().replace(/---+$/, "").trim();
-      values[state.currField.name] = validateAndParseFieldValue(
-        state.currField,
-        val
-      );
+    if (xstate.currField) {
+      const val = content.substring(xstate.s, e);
+      const parsedValue = validateAndParseFieldValue(xstate.currField, val);
+      if (parsedValue !== void 0) {
+        values[xstate.currField.name] = parsedValue;
+      }
+    }
+    checkMissingRequiredFields(xstate, values, index);
+    xstate.s = e + prefix.length;
+    xstate.currField = field;
+    if (!xstate.extractedFields.includes(field)) {
+      xstate.extractedFields.push(field);
     }
-    state.s = e + prefix.length;
-    state.currField = field;
   }
 };
-var streamingExtractFinalValue = (values, state, content) => {
-  if (!state.currField) {
-    return;
+var streamingExtractFinalValue = (sig, values, xstate, content) => {
+  if (xstate.currField) {
+    const val = content.substring(xstate.s);
+    const parsedValue = validateAndParseFieldValue(xstate.currField, val);
+    if (parsedValue !== void 0) {
+      values[xstate.currField.name] = parsedValue;
+    }
   }
-  const val = content.substring(state.s).trim().replace(/---+$/, "").trim();
-  values[state.currField.name] = validateAndParseFieldValue(
-    state.currField,
-    val
-  );
+  const fields = sig.getOutputFields();
+  checkMissingRequiredFields(xstate, values, fields.length - 1);
 };
 var convertValueToType = (field, val) => {
   switch (field.type?.name) {
@@ -4692,59 +4787,98 @@ var convertValueToType = (field, val) => {
       return 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 ValidationError({ message, field, value });
-};
+function* streamingValues(sig, values, xstate, content) {
+  if (!xstate.currField) {
+    return;
+  }
+  const fieldName = xstate.currField.name;
+  if (!xstate.streamedIndex) {
+    xstate.streamedIndex = { [fieldName]: 0 };
+  }
+  if (!xstate.currField.type || !xstate.currField.type.isArray && xstate.currField.type.name === "string") {
+    const s = xstate.s + (xstate.streamedIndex[fieldName] ?? 0);
+    const v = content.substring(s);
+    yield { [fieldName]: v };
+    xstate.streamedIndex[fieldName] = v.length;
+    return;
+  }
+  for (const key of Object.keys(values)) {
+    const value = values[key];
+    if (Array.isArray(value)) {
+      const s = xstate.streamedIndex[fieldName] ?? 0;
+      const v = value.slice(s);
+      if (v) {
+        yield { [fieldName]: v };
+        xstate.streamedIndex[fieldName] = s + 1;
+      }
+      continue;
+    }
+    if (!xstate.streamedIndex[fieldName]) {
+      yield { [fieldName]: value };
+      xstate.streamedIndex[fieldName] = 1;
+    }
+  }
+}
 function validateAndParseFieldValue(field, fieldValue) {
-  const fv = fieldValue?.toLocaleLowerCase();
-  if (!fieldValue || !fv || fv === "" || fv === "null" || fv === "undefined") {
+  const fv = fieldValue?.trim();
+  if (!fv || !fv || fv === "" || fv === "null" || fv === "NULL" || fv === "undefined") {
     if (field.isOptional) {
       return;
     }
-    throw expectedTypeError(field, new Error("Empty value"), fieldValue);
+    throw new ValidationError({
+      message: "Required field is missing",
+      fields: [field],
+      value: fv
+    });
   }
-  let value = fieldValue;
+  let value;
   if (field.type?.name === "json") {
     try {
-      const text = extractBlock(fieldValue);
+      const text = extractBlock(fv);
       value = JSON5.parse(text);
       return value;
     } catch (e) {
-      throw expectedTypeError(field, e, fieldValue);
+      throw new ValidationError({
+        message: "Invalid JSON: " + e.message,
+        fields: [field],
+        value: fv
+      });
     }
   }
   if (field.type?.isArray) {
     try {
       try {
-        value = JSON5.parse(fieldValue);
+        value = JSON5.parse(fv);
       } catch {
-        value = parseMarkdownList(fieldValue);
+        value = parseMarkdownList(fv);
       }
       if (!Array.isArray(value)) {
         throw new Error("Expected an array");
       }
     } catch (e) {
-      throw expectedTypeError(field, e, fieldValue);
+      throw new ValidationError({
+        message: "Invalid array: " + e.message,
+        fields: [field],
+        value: fv
+      });
     }
   }
-  if (Array.isArray(value)) {
-    for (const [index, item] of value.entries()) {
-      try {
+  try {
+    if (Array.isArray(value)) {
+      for (const [index, item] of value.entries()) {
         value[index] = convertValueToType(field, item);
-      } catch (e) {
-        throw expectedTypeError(field, e, item);
       }
+    } else {
+      value = convertValueToType(field, fv);
     }
-  } else {
-    try {
-      value = convertValueToType(field, fieldValue);
-    } catch (e) {
-      throw expectedTypeError(field, e, fieldValue);
-    }
+  } catch (e) {
+    throw new ValidationError({
+      message: e.message,
+      fields: [field],
+      value: fieldValue
+    });
   }
-  return value;
+  return value ?? fv;
 }
 var extractBlock = (input) => {
   const jsonBlockPattern = /```([A-Za-z]+)?\s*([\s\S]*?)\s*```/g;
@@ -4911,7 +5045,6 @@ function parseFunctionCalls(ai, functionCalls, values, model) {
 }
 
 // dsp/generate.ts
-var colorLog4 = new ColorLog();
 var AxGen = class extends AxProgramWithSignature {
   promptTemplate;
   asserts;
@@ -4976,7 +5109,7 @@ var AxGen = class extends AxProgramWithSignature {
     );
     return res;
   }
-  async forwardCore({
+  async *forwardCore({
     ai,
     mem,
     options
@@ -4992,7 +5125,18 @@ var AxGen = class extends AxProgramWithSignature {
       options
     });
     if (res instanceof ReadableStream) {
-      return await this.processSteamingResponse({
+      yield* this.processStreamingResponse({
+        ai,
+        model,
+        res,
+        usageInfo,
+        mem,
+        traceId,
+        sessionId,
+        functions
+      });
+    } else {
+      yield await this.processResponse({
         ai,
         model,
         res,
@@ -5003,18 +5147,8 @@ var AxGen = class extends AxProgramWithSignature {
         functions
       });
     }
-    return await this.processResponse({
-      ai,
-      model,
-      res,
-      usageInfo,
-      mem,
-      traceId,
-      sessionId,
-      functions
-    });
   }
-  async processSteamingResponse({
+  async *processStreamingResponse({
     ai,
     model,
     res,
@@ -5026,36 +5160,39 @@ var AxGen = class extends AxProgramWithSignature {
   }) {
     const functionCalls = [];
     const values = {};
-    const xstate = { s: -1 };
+    const xstate = { extractedFields: [], s: -1 };
     let content = "";
     for await (const v of res) {
-      for (const result of v.results ?? []) {
-        if (v.modelUsage) {
-          this.usage.push({ ...usageInfo, ...v.modelUsage });
-        }
-        if (result.content) {
-          content += result.content;
-          mem.updateResult({ name: result.name, content }, sessionId);
-          assertStreamingAssertions(
-            this.streamingAsserts,
-            values,
-            xstate,
-            content,
-            false
-          );
-          streamingExtractValues(this.signature, values, xstate, content);
-          assertAssertions(this.asserts, values);
-        }
-        if (result.functionCalls) {
-          mergeFunctionCalls(functionCalls, result.functionCalls);
-          mem.updateResult(
-            { name: result.name, content, functionCalls },
-            sessionId
-          );
-        }
-        if (result.finishReason === "length") {
-          throw new Error("Max tokens reached before completion");
-        }
+      const result = v.results[0];
+      if (!result) {
+        continue;
+      }
+      if (v.modelUsage) {
+        this.usage.push({ ...usageInfo, ...v.modelUsage });
+      }
+      if (result.content) {
+        content += result.content;
+        mem.updateResult({ name: result.name, content }, sessionId);
+        assertStreamingAssertions(
+          this.streamingAsserts,
+          values,
+          xstate,
+          content,
+          false
+        );
+        streamingExtractValues(this.signature, values, xstate, content);
+        assertAssertions(this.asserts, values);
+        yield* streamingValues(this.signature, values, xstate, content);
+      }
+      if (result.functionCalls) {
+        mergeFunctionCalls(functionCalls, result.functionCalls);
+        mem.updateResult(
+          { name: result.name, content, functionCalls },
+          sessionId
+        );
+      }
+      if (result.finishReason === "length") {
+        throw new Error("Max tokens reached before completion");
       }
     }
     const funcs = parseFunctionCalls(ai, functionCalls, values, model);
@@ -5073,7 +5210,6 @@ var AxGen = class extends AxProgramWithSignature {
       );
       this.functionsExecuted = /* @__PURE__ */ new Set([...this.functionsExecuted, ...fx]);
     }
-    streamingExtractFinalValue(values, xstate, content);
     assertStreamingAssertions(
       this.streamingAsserts,
       values,
@@ -5081,8 +5217,9 @@ var AxGen = class extends AxProgramWithSignature {
       content,
       true
     );
+    streamingExtractFinalValue(this.signature, values, xstate, content);
     assertAssertions(this.asserts, values);
-    return { ...values };
+    return values;
   }
   async processResponse({
     ai,
@@ -5094,45 +5231,47 @@ var AxGen = class extends AxProgramWithSignature {
     functions
   }) {
     const values = {};
-    for (const result of res.results ?? []) {
-      if (res.modelUsage) {
-        this.usage.push({ ...usageInfo, ...res.modelUsage });
-      }
-      mem.addResult(result, sessionId);
-      if (result.content) {
-        extractValues(this.signature, values, result.content);
-        assertAssertions(this.asserts, values);
-      }
-      if (result.functionCalls) {
-        const funcs = parseFunctionCalls(ai, result.functionCalls, values);
-        if (funcs) {
-          if (!functions) {
-            throw new Error("Functions are not defined");
-          }
-          const fx = await processFunctions(
-            ai,
-            functions,
-            funcs,
-            mem,
-            sessionId,
-            traceId
-          );
-          this.functionsExecuted = /* @__PURE__ */ new Set([...this.functionsExecuted, ...fx]);
+    const result = res.results[0];
+    if (!result) {
+      throw new Error("No result found");
+    }
+    if (res.modelUsage) {
+      this.usage.push({ ...usageInfo, ...res.modelUsage });
+    }
+    mem.addResult(result, sessionId);
+    if (result.content) {
+      extractValues(this.signature, values, result.content);
+      assertAssertions(this.asserts, values);
+    }
+    if (result.functionCalls) {
+      const funcs = parseFunctionCalls(ai, result.functionCalls, values);
+      if (funcs) {
+        if (!functions) {
+          throw new Error("Functions are not defined");
         }
-      }
-      if (result.finishReason === "length") {
-        throw new Error("Max tokens reached before completion");
+        const fx = await processFunctions(
+          ai,
+          functions,
+          funcs,
+          mem,
+          sessionId,
+          traceId
+        );
+        this.functionsExecuted = /* @__PURE__ */ new Set([...this.functionsExecuted, ...fx]);
       }
     }
+    if (result.finishReason === "length") {
+      throw new Error("Max tokens reached before completion");
+    }
     return { ...values };
   }
-  async _forward(ai, values, options, span) {
+  async *_forward2(ai, values, options, span) {
     const stopFunction = (options?.stopFunction ?? this.options?.stopFunction)?.toLowerCase();
-    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();
+    const maxRetries = options.maxRetries ?? this.options?.maxRetries ?? 10;
+    const maxSteps = options.maxSteps ?? this.options?.maxSteps ?? 10;
+    const mem = options.mem ?? this.options?.mem ?? new AxMemory();
     let err;
-    if (options?.functions && options?.functions.length > 0) {
+    if (options?.functions && options.functions.length > 0) {
       const promptTemplate = this.options?.promptTemplate ?? AxPromptTemplate;
       this.promptTemplate = new promptTemplate(
         this.signature,
@@ -5147,25 +5286,22 @@ var AxGen = class extends AxProgramWithSignature {
     multiStepLoop: for (let n = 0; n < maxSteps; n++) {
       for (let errCount = 0; errCount < maxRetries; errCount++) {
         try {
-          const output = await this.forwardCore({
-            options,
-            ai,
-            mem
-          });
+          const generator = this.forwardCore({ options, ai, mem });
+          for await (const delta of generator) {
+            yield {
+              version: errCount,
+              delta
+            };
+          }
           const lastMemItem = mem.getLast(options?.sessionId);
-          if (lastMemItem) {
-            const stopFunctionExecuted = stopFunction && this.functionsExecuted.has(stopFunction);
-            if (lastMemItem.role === "function") {
-              if (!stopFunction || !stopFunctionExecuted) {
-                continue multiStepLoop;
-              }
-            }
-            if (!stopFunctionExecuted) {
-              assertRequiredFields(this.signature, output);
-            }
+          const shouldContinue = this.shouldContinueSteps(
+            lastMemItem,
+            stopFunction
+          );
+          if (shouldContinue) {
+            continue multiStepLoop;
           }
-          this.trace = { ...values, ...output };
-          return output;
+          return;
         } catch (e) {
           let errorFields;
           span?.recordException(e);
@@ -5180,23 +5316,13 @@ var AxGen = class extends AxProgramWithSignature {
             throw e;
           }
           if (errorFields) {
-            mem.add(
-              {
-                role: "user",
-                content: this.promptTemplate.renderExtraFields(errorFields)
-              },
-              options?.sessionId
+            handleValidationError(
+              mem,
+              errorFields,
+              ai,
+              this.promptTemplate,
+              options.sessionId
             );
-            mem.addTag("error");
-            if (ai.getOptions().debug) {
-              process.stdout.write(
-                colorLog4.red(
-                  `Error Correction:
-${JSON.stringify(errorFields, null, 2)}
-`
-                )
-              );
-            }
           }
         }
       }
@@ -5207,35 +5333,68 @@ ${JSON.stringify(errorFields, null, 2)}
     }
     throw new Error(`Max steps reached: ${maxSteps}`);
   }
-  async forward(ai, values, options) {
+  shouldContinueSteps(lastMemItem, stopFunction) {
+    const stopFunctionExecuted = stopFunction && this.functionsExecuted.has(stopFunction);
+    if (lastMemItem?.role === "function" && stopFunction && stopFunctionExecuted) {
+      return false;
+    }
+    if (lastMemItem?.role === "function") {
+      return true;
+    }
+    return false;
+  }
+  async *_forward1(ai, values, options) {
     const tracer = this.options?.tracer ?? options?.tracer;
     let functions = this.functions;
     if (options?.functions) {
       functions = parseFunctions(options.functions, this.functions);
     }
     if (!tracer) {
-      return await this._forward(ai, values, {
+      yield* this._forward2(ai, values, {
         ...options,
         functions
       });
+      return;
     }
     const funcNames = functions?.map((f) => f.name).join(",");
     const attributes = {
       ["generate.signature"]: this.signature.toString(),
       ["generate.functions"]: funcNames ?? ""
     };
-    return await tracer.startActiveSpan(
-      "Generate",
-      {
-        kind: SpanKind2.SERVER,
-        attributes
-      },
-      async (span) => {
-        const res = this._forward(ai, values, options, span);
-        span.end();
-        return res;
-      }
-    );
+    const span = tracer.startSpan("Generate", {
+      kind: SpanKind2.SERVER,
+      attributes
+    });
+    try {
+      yield* this._forward2(
+        ai,
+        values,
+        {
+          ...options,
+          functions
+        },
+        span
+      );
+    } finally {
+      span.end();
+    }
+  }
+  async forward(ai, values, options) {
+    const generator = this._forward1(ai, values, {
+      ...options
+    });
+    let buffer = {};
+    for await (const delta of generator) {
+      buffer = mergeDeltas(buffer, delta.delta);
+    }
+    this.trace = { ...values, ...buffer };
+    return buffer;
+  }
+  async *streamingForward(ai, values, options) {
+    yield* this._forward1(ai, values, {
+      ...options,
+      stream: true
+    });
   }
 };
 
@@ -5325,7 +5484,7 @@ var AxAgent = class {
       func: wrappedFunc
     };
   }
-  async forward(ai, values, options) {
+  init(ai, options) {
     const _ai = this.ai ?? ai;
     const funcs = [
       ...options?.functions ?? [],
@@ -5336,26 +5495,15 @@ var AxAgent = class {
       const opt2 = { ...options, functions: funcs };
       this.program = new AxGen(this.signature, opt2);
     }
-    if (!options?.tracer) {
-      return await this.program.forward(_ai, values, opt);
-    }
-    const attributes = {
-      ["agent.name"]: this.name,
-      ["agent.description"]: this.description,
-      ["agent.subAgents"]: this.subAgentList ?? "none"
-    };
-    return await options?.tracer.startActiveSpan(
-      "Agent",
-      {
-        kind: SpanKind3.SERVER,
-        attributes
-      },
-      async (span) => {
-        const res = await this.program.forward(_ai, values, opt);
-        span.end();
-        return res;
-      }
-    );
+    return { _ai, opt };
+  }
+  async forward(ai, values, options) {
+    const { _ai, opt } = this.init(ai, options);
+    return await this.program.forward(_ai, values, opt);
+  }
+  async *streamingForward(ai, values, options) {
+    const { _ai, opt } = this.init(ai, options);
+    return yield* this.program.streamingForward(_ai, values, opt);
   }
 };
 function toCamelCase(inputString) {
@@ -5626,7 +5774,7 @@ var randomSample = (array, n) => {
 };
 
 // db/base.ts
-import { SpanKind as SpanKind4 } from "@opentelemetry/api";
+import { SpanKind as SpanKind3 } from "@opentelemetry/api";
 var AxDBBase = class {
   name;
   fetch;
@@ -5653,7 +5801,7 @@ var AxDBBase = class {
     return await this.tracer?.startActiveSpan(
       "DB Upsert Request",
       {
-        kind: SpanKind4.SERVER,
+        kind: SpanKind3.SERVER,
         attributes: {
           [axSpanAttributes.DB_SYSTEM]: this.name,
           [axSpanAttributes.DB_OPERATION_NAME]: "upsert",
@@ -5663,9 +5811,11 @@ var AxDBBase = class {
         }
       },
       async (span) => {
-        const res = await this._upsert(req, update, { span });
-        span.end();
-        return res;
+        try {
+          return await this._upsert(req, update, { span });
+        } finally {
+          span.end();
+        }
       }
     );
   }
@@ -5685,7 +5835,7 @@ var AxDBBase = class {
     return await this.tracer?.startActiveSpan(
       "DB Batch Upsert Request",
       {
-        kind: SpanKind4.SERVER,
+        kind: SpanKind3.SERVER,
         attributes: {
           [axSpanAttributes.DB_SYSTEM]: this.name,
           [axSpanAttributes.DB_OPERATION_NAME]: "upsert",
@@ -5695,9 +5845,11 @@ var AxDBBase = class {
         }
       },
       async (span) => {
-        const res = await this._batchUpsert(req, update, { span });
-        span.end();
-        return res;
+        try {
+          return await this._batchUpsert(req, update, { span });
+        } finally {
+          span.end();
+        }
       }
     );
   }
@@ -5711,7 +5863,7 @@ var AxDBBase = class {
     return await this.tracer?.startActiveSpan(
       "DB Query Request",
       {
-        kind: SpanKind4.SERVER,
+        kind: SpanKind3.SERVER,
         attributes: {
           [axSpanAttributes.DB_SYSTEM]: this.name,
           [axSpanAttributes.DB_OPERATION_NAME]: "upsert",
@@ -5721,9 +5873,11 @@ var AxDBBase = class {
         }
       },
       async (span) => {
-        const res = await this._query(req, { span });
-        span.end();
-        return res;
+        try {
+          return await this._query(req, { span });
+        } finally {
+          span.end();
+        }
       }
     );
   }
@@ -6990,6 +7144,122 @@ var AxEmbeddingAdapter = class {
   }
 };
 
+// ai/mock/api.ts
+var AxMockAIService = class {
+  constructor(config = {}) {
+    this.config = config;
+  }
+  options = {};
+  metrics = {
+    latency: {
+      chat: { mean: 0, p95: 0, p99: 0, samples: [] },
+      embed: { mean: 0, p95: 0, p99: 0, samples: [] }
+    },
+    errors: {
+      chat: { count: 0, rate: 0, total: 0 },
+      embed: { count: 0, rate: 0, total: 0 }
+    }
+  };
+  getName() {
+    return this.config.name ?? "mock-ai-service";
+  }
+  getModelInfo() {
+    return {
+      name: "mock-model",
+      provider: "mock-provider",
+      promptTokenCostPer1M: 100,
+      completionTokenCostPer1M: 100,
+      ...this.config.modelInfo
+    };
+  }
+  getEmbedModelInfo() {
+    return this.config.embedModelInfo;
+  }
+  // eslint-disable-next-line @typescript-eslint/no-unused-vars
+  getFeatures(_model) {
+    return {
+      functions: this.config.features?.functions ?? false,
+      streaming: this.config.features?.streaming ?? false
+    };
+  }
+  getModelMap() {
+    return this.config.modelMap;
+  }
+  getMetrics() {
+    return this.metrics;
+  }
+  async chat(req, _options) {
+    if (this.config.latencyMs) {
+      await new Promise((resolve) => setTimeout(resolve, this.config.latencyMs));
+    }
+    if (this.config.shouldError) {
+      throw new Error(this.config.errorMessage ?? "Mock chat error");
+    }
+    this.updateMetrics("chat");
+    if (typeof this.config.chatResponse === "function") {
+      return this.config.chatResponse(req);
+    }
+    return this.config.chatResponse ?? {
+      results: [
+        {
+          content: "Mock response",
+          finishReason: "stop"
+        }
+      ],
+      modelUsage: {
+        promptTokens: 10,
+        completionTokens: 5,
+        totalTokens: 15
+      }
+    };
+  }
+  async embed(req, _options) {
+    if (this.config.latencyMs) {
+      await new Promise((resolve) => setTimeout(resolve, this.config.latencyMs));
+    }
+    if (this.config.shouldError) {
+      throw new Error(this.config.errorMessage ?? "Mock embed error");
+    }
+    this.updateMetrics("embed");
+    if (typeof this.config.embedResponse === "function") {
+      return this.config.embedResponse(req);
+    }
+    return this.config.embedResponse ?? {
+      embeddings: [[0.1, 0.2, 0.3]],
+      modelUsage: {
+        promptTokens: 5,
+        completionTokens: 0,
+        totalTokens: 5
+      }
+    };
+  }
+  setOptions(options) {
+    this.options = options;
+  }
+  getOptions() {
+    return this.options;
+  }
+  updateMetrics(type) {
+    const latency = this.config.latencyMs ?? 0;
+    this.metrics.latency[type].samples.push(latency);
+    const samples = this.metrics.latency[type].samples;
+    this.metrics.latency[type].mean = samples.reduce((a, b) => a + b, 0) / samples.length;
+    if (samples.length > 0) {
+      const sortedSamples = [...samples].sort((a, b) => a - b);
+      const p95Index = Math.max(0, Math.floor(sortedSamples.length * 0.95) - 1);
+      this.metrics.latency[type].p95 = sortedSamples[p95Index] ?? latency;
+      const p99Index = Math.max(0, Math.floor(sortedSamples.length * 0.99) - 1);
+      this.metrics.latency[type].p99 = sortedSamples[p99Index] ?? latency;
+    }
+    if (this.config.shouldError) {
+      this.metrics.errors[type].count++;
+      this.metrics.errors[type].total++;
+      const totalRequests = this.metrics.latency[type].samples.length;
+      this.metrics.errors[type].rate = totalRequests > 0 ? this.metrics.errors[type].count / totalRequests : 0;
+    }
+  }
+};
+
 // prompts/rag.ts
 var AxRAG = class extends AxChainOfThought {
   genQuery;
@@ -7078,6 +7348,7 @@ export {
   AxJSInterpreterPermission,
   AxLLMRequestTypeValues,
   AxMemory,
+  AxMockAIService,
   AxProgram,
   AxProgramWithSignature,
   AxPromptTemplate,