@ax-llm/ax 11.0.18 → 11.0.20

package/index.js

@@ -3985,7 +3985,10 @@ var SignatureParser = class {
       this.skipWhitespace();
       const optionalDesc = this.parseParsedString();
       this.skipWhitespace();
-      const inputs = this.parseFieldList(this.parseField.bind(this), "input");
+      const inputs = this.parseFieldList(
+        this.parseInputField.bind(this),
+        "input"
+      );
       this.skipWhitespace();
       if (this.position >= this.input.length) {
         throw new Error(
@@ -3999,7 +4002,10 @@ var SignatureParser = class {
           'Incomplete signature: No output fields specified after "->"'
         );
       }
-      const outputs = this.parseFieldList(this.parseField.bind(this), "output");
+      const outputs = this.parseFieldList(
+        this.parseOutputField.bind(this),
+        "output"
+      );
       return {
         desc: optionalDesc?.trim(),
         inputs,
@@ -4061,11 +4067,75 @@ ${pointer}`;
     }
     return fields;
   }
-  parseField() {
+  // -------------------------------
+  // Parse input fields (no "class" type and no internal flag)
+  // -------------------------------
+  parseInputField() {
     this.skipWhitespace();
     const name = this.parseParsedIdentifier();
     this.currentFieldName = name;
-    const isOptional = this.match("?");
+    let isOptional = void 0;
+    while (true) {
+      if (this.match("?")) {
+        isOptional = true;
+        continue;
+      }
+      if (this.match("!")) {
+        throw new Error(
+          `Input field "${name}" does not support the internal marker "!"`
+        );
+      }
+      break;
+    }
+    let type;
+    this.skipWhitespace();
+    if (this.match(":")) {
+      this.skipWhitespace();
+      if (/^class\b/.test(this.input.slice(this.position))) {
+        throw new Error(
+          `Input field "${name}" does not support the "class" type`
+        );
+      } else {
+        try {
+          const typeName = this.parseTypeNotClass();
+          const isArray = this.match("[]");
+          type = { name: typeName, isArray };
+        } catch (error) {
+          throw new Error(
+            `Input field "${name}": ${error instanceof Error ? error.message : "Unknown error"}`
+          );
+        }
+      }
+    }
+    this.skipWhitespace();
+    const desc = this.parseParsedString();
+    return {
+      name,
+      desc: desc?.trim(),
+      type,
+      isOptional
+    };
+  }
+  // -------------------------------
+  // Parse output fields (supports both "class" type and the internal marker)
+  // -------------------------------
+  parseOutputField() {
+    this.skipWhitespace();
+    const name = this.parseParsedIdentifier();
+    this.currentFieldName = name;
+    let isOptional = false;
+    let isInternal = false;
+    while (true) {
+      if (this.match("?")) {
+        isOptional = true;
+        continue;
+      }
+      if (this.match("!")) {
+        isInternal = true;
+        continue;
+      }
+      break;
+    }
     let type;
     this.skipWhitespace();
     if (this.match(":")) {
@@ -4073,16 +4143,16 @@ ${pointer}`;
       if (this.match("class")) {
         const isArray = this.match("[]");
         this.skipWhitespace();
-        const desc2 = this.parseParsedString();
-        if (!desc2) {
+        const classNamesString = this.parseParsedString();
+        if (!classNamesString) {
           throw new Error(
-            `Field "${name}": Expected class names in quotes after "class" type. Example: class "MyClass1, MyClass2"`
+            `Output field "${name}": Expected class names in quotes after "class" type. Example: class "MyClass1, MyClass2"`
           );
         }
-        const classes = desc2.split(/[,\s]+/).map((s) => s.trim()).filter((s) => s.length > 0);
+        const classes = classNamesString.split(/[,\s]+/).map((s) => s.trim()).filter((s) => s.length > 0);
         if (classes.length === 0) {
           throw new Error(
-            `Field "${name}": Empty class list provided. At least one class name is required`
+            `Output field "${name}": Empty class list provided. At least one class name is required`
           );
         }
         type = { name: "class", isArray, classes };
@@ -4093,28 +4163,20 @@ ${pointer}`;
           type = { name: typeName, isArray };
         } catch (error) {
           throw new Error(
-            `Field "${name}": ${error instanceof Error ? error.message : "Unknown error"}`
+            `Output field "${name}": ${error instanceof Error ? error.message : "Unknown error"}`
           );
         }
       }
     }
     this.skipWhitespace();
     const desc = this.parseParsedString();
-    if (type?.name === "class") {
-      return {
-        name,
-        desc: desc?.trim(),
-        type,
-        isOptional
-      };
-    } else {
-      return {
-        name,
-        desc: desc?.trim(),
-        type,
-        isOptional
-      };
-    }
+    return {
+      name,
+      desc: desc?.trim(),
+      type,
+      isOptional,
+      isInternal
+    };
   }
   parseTypeNotClass() {
     const types = [
@@ -4268,8 +4330,9 @@ var AxSignature = class _AxSignature {
       name: field.name,
       title,
       description: "desc" in field ? field.desc : void 0,
-      isOptional: field.isOptional,
-      type: field.type ?? { name: "string", isArray: false }
+      type: field.type ?? { name: "string", isArray: false },
+      ..."isInternal" in field ? { isInternal: field.isInternal } : {},
+      ..."isOptional" in field ? { isOptional: field.isOptional } : {}
     };
   };
   parseField = (field) => {
@@ -4303,7 +4366,7 @@ var AxSignature = class _AxSignature {
   getOutputFields = () => this.outputFields;
   getDescription = () => this.description;
   toTitle = (name) => {
-    let result = name.replaceAll("_", " ");
+    let result = name.replace(/_/g, " ");
     result = result.replace(/([A-Z]|[0-9]+)/g, " $1").trim();
     return result.charAt(0).toUpperCase() + result.slice(1);
   };
@@ -5117,7 +5180,7 @@ var renderOutputFields = (fields) => {
   const rows = fields.map((field) => {
     const name = field.title;
     const type = field.type?.name ? toFieldType(field.type) : "string";
-    const requiredMsg = field.isOptional ? `Only include this ${type} field is it's value is available` : `This ${type} field must be included`;
+    const requiredMsg = field.isOptional ? `Only include this ${type} field if its value is available` : `This ${type} field must be included`;
     const description = field.description ? ` ${formatDescription(field.description)}` : "";
     return `${name}: (${requiredMsg})${description}`.trim();
   });
@@ -5305,7 +5368,7 @@ var formatDateWithTimezone = (date) => {
 
 // dsp/extract.ts
 var extractValues = (sig, values, content) => {
-  const xstate = { extractedFields: [], s: -1 };
+  const xstate = { extractedFields: [], streamedIndex: {}, s: -1 };
   streamingExtractValues(sig, values, xstate, content);
   streamingExtractFinalValue(sig, values, xstate, content);
 };
@@ -5359,7 +5422,7 @@ var streamingExtractValues = (sig, values, xstate, content, streamingValidation
       if (parsedValue !== void 0) {
         values[xstate.currField.name] = parsedValue;
       }
-      xstate.lastDelta = val;
+      xstate.prevField = { field: xstate.currField, s: xstate.s, e };
     }
     checkMissingRequiredFields(xstate, values, index);
     xstate.s = e + prefixLen;
@@ -5368,6 +5431,9 @@ var streamingExtractValues = (sig, values, xstate, content, streamingValidation
     if (!xstate.extractedFields.includes(field)) {
       xstate.extractedFields.push(field);
     }
+    if (xstate.streamedIndex[field.name] === void 0) {
+      xstate.streamedIndex[field.name] = 0;
+    }
   }
 };
 var streamingExtractFinalValue = (sig, values, xstate, content) => {
@@ -5423,25 +5489,18 @@ var convertValueToType = (field, val) => {
       return val;
   }
 };
-function processStreamingDelta(content, xstate) {
-  if (!xstate.currField) {
-    return null;
-  }
-  const { name: fieldName, type: fieldType } = xstate.currField ?? {};
-  const { isArray: fieldIsArray, name: fieldTypeName } = fieldType ?? {};
-  if (!xstate.streamedIndex) {
-    xstate.streamedIndex = { [fieldName]: 0 };
-  }
-  if (fieldIsArray) {
-    return null;
-  }
-  if (fieldTypeName !== "string" && fieldTypeName !== "code") {
-    return null;
+function* yieldDelta(content, field, s, e, xstate) {
+  const { name: fieldName, isInternal } = field;
+  const { isArray: fieldIsArray, name: fieldTypeName } = field.type ?? {};
+  if (isInternal || fieldIsArray || fieldTypeName && fieldTypeName !== "string" && fieldTypeName !== "code") {
+    return;
   }
   const pos = xstate.streamedIndex[fieldName] ?? 0;
-  const s = xstate.s + pos;
   const isFirstChunk = pos === 0;
-  const d1 = content.substring(s);
+  const d1 = content.substring(s + pos, e);
+  if (d1.length === 0) {
+    return;
+  }
   let d2 = d1.replace(/\s+$/, "");
   if (xstate.currField?.type?.name === "code") {
     d2 = d2.replace(/\s*```\s*$/, "");
@@ -5451,41 +5510,45 @@ function processStreamingDelta(content, xstate) {
     d3 = d3.replace(/^[ ]*```[a-zA-Z0-9]*\n\s*/, "");
   }
   if (d3.length > 0) {
+    yield { [fieldName]: d3 };
     xstate.streamedIndex[fieldName] = pos + d2.length;
   }
-  return d3;
 }
-function getStreamingDelta(content, xstate) {
-  if (xstate.lastDelta) {
-    processStreamingDelta(xstate.lastDelta, xstate);
-    xstate.lastDelta = void 0;
+function* streamValues(sig, content, values, xstate) {
+  if (xstate.prevField && !xstate.prevField.field.isInternal) {
+    const { field, s, e } = xstate.prevField;
+    yield* yieldDelta(content, field, s, e, xstate);
+    xstate.prevField = void 0;
   }
-  return processStreamingDelta(content, xstate);
-}
-function* streamValues(sig, values, xstate, delta) {
-  if (!xstate.currField) {
-    return;
-  }
-  const fieldName = xstate.currField.name;
-  if (delta && delta.length > 0) {
-    yield { [fieldName]: delta };
+  if (!xstate.currField || xstate.currField.isInternal) {
     return;
   }
+  yield* yieldDelta(
+    content,
+    xstate.currField,
+    xstate.s,
+    content.length,
+    xstate
+  );
+  const outputFields = sig.getOutputFields();
   for (const key of Object.keys(values)) {
+    const field = outputFields.find((f) => f.name === key);
+    if (!field || field.isInternal) {
+      continue;
+    }
     const value = values[key];
     if (Array.isArray(value)) {
       const s = xstate.streamedIndex?.[key] ?? 0;
       const v = value.slice(s);
       if (v && v.length > 0) {
         yield { [key]: v };
-        if (xstate.streamedIndex) {
-          xstate.streamedIndex[key] = s + 1;
-        } else {
-          xstate.streamedIndex = { [key]: s + 1 };
-        }
+        xstate.streamedIndex[key] = s + 1;
       }
-    } else {
+      continue;
+    }
+    if (!xstate.streamedIndex[key]) {
       yield { [key]: value };
+      xstate.streamedIndex[key] = 1;
     }
   }
 }
@@ -5977,6 +6040,7 @@ var AxGen = class extends AxProgramWithSignature {
     const values = {};
     const xstate = {
       extractedFields: [],
+      streamedIndex: {},
       s: -1
     };
     let content = "";
@@ -6012,11 +6076,7 @@ var AxGen = class extends AxProgramWithSignature {
           content,
           streamingValidation
         );
-        if (skip) {
-          continue;
-        }
         assertStreamingAssertions(this.streamingAsserts, xstate, content);
-        assertAssertions(this.asserts, values);
         if (this.streamingFieldProcessors.length !== 0) {
           await processStreamingFieldProcessors(
             this.streamingFieldProcessors,
@@ -6027,8 +6087,11 @@ var AxGen = class extends AxProgramWithSignature {
             sessionId
           );
         }
-        const delta = getStreamingDelta(content, xstate);
-        yield* streamValues(this.signature, values, xstate, delta);
+        yield* streamValues(this.signature, content, values, xstate);
+        if (skip) {
+          continue;
+        }
+        assertAssertions(this.asserts, values);
       }
       if (result.finishReason === "length") {
         throw new Error("Max tokens reached before completion");
@@ -6071,8 +6134,7 @@ var AxGen = class extends AxProgramWithSignature {
           true
         );
       }
-      const delta = getStreamingDelta(content, xstate);
-      yield* streamValues(this.signature, values, xstate, delta);
+      yield* streamValues(this.signature, content, values, xstate);
     }
   }
   async processResponse({
@@ -6126,6 +6188,12 @@ var AxGen = class extends AxProgramWithSignature {
         throw new Error("Max tokens reached before completion");
       }
     }
+    const publicValues = { ...values };
+    for (const field of this.signature.getOutputFields()) {
+      if (field.isInternal) {
+        delete publicValues[field.name];
+      }
+    }
     return { ...values };
   }
   async *_forward2(ai, values, options, span) {
@@ -8010,6 +8078,125 @@ var AxJSInterpreter = class {
   }
 };
 
+// ai/mock/api.ts
+var AxMockAIService = class {
+  constructor(config = {}) {
+    this.config = config;
+    this.config.id = this.config.id ?? crypto.randomUUID();
+  }
+  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";
+  }
+  getId() {
+    return this.config.id ?? "mock-ai-service-id";
+  }
+  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
+    };
+  }
+  getModelList() {
+    return this.config.models;
+  }
+  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 await 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.config.options = options;
+  }
+  getOptions() {
+    return this.config.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;
+    }
+  }
+};
+
 // dsp/router.ts
 var colorLog6 = new ColorLog();
 var AxRoute = class {
@@ -8279,126 +8466,6 @@ var AxEmbeddingAdapter = class {
   }
 };
 
-// ai/mock/api.ts
-var AxMockAIService = class {
-  constructor(config = {}) {
-    this.config = config;
-    this.config.id = this.config.id ?? crypto.randomUUID();
-  }
-  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";
-  }
-  getId() {
-    return this.config.id ?? "mock-ai-service-id";
-  }
-  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
-    };
-  }
-  getModelList() {
-    return this.config.models;
-  }
-  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;