@ax-llm/ax 10.0.40 → 10.0.41

package/index.js

@@ -789,6 +789,61 @@ var setResponseAttr = (res, span) => {
   }
 };
 
+// ai/google-vertex/auth.ts
+import { GoogleAuth } from "google-auth-library";
+var GoogleVertexAuth = class {
+  auth;
+  client;
+  currentToken;
+  tokenExpiry;
+  constructor(config = {}) {
+    this.auth = new GoogleAuth({
+      scopes: ["https://www.googleapis.com/auth/cloud-platform"],
+      ...config
+    });
+  }
+  async getAuthenticatedClient() {
+    if (!this.client) {
+      this.client = await this.auth.getClient();
+    }
+    return this.client;
+  }
+  async getAccessToken() {
+    if (this.currentToken && this.tokenExpiry && Date.now() < this.tokenExpiry) {
+      return this.currentToken;
+    }
+    const client = await this.getAuthenticatedClient();
+    const tokenResponse = await client.getAccessToken();
+    this.currentToken = tokenResponse.token ?? void 0;
+    const expiry = this.getExpiry(tokenResponse);
+    const fiveMinutes = 5 * 60 * 1e3;
+    this.tokenExpiry = expiry - fiveMinutes;
+    return this.currentToken;
+  }
+  /**
+   * Get the expiry date from the token response.
+   */
+  getExpiry(tokenResponse) {
+    const oneHour = 3600 * 1e3;
+    let expiry = Date.now() + oneHour;
+    let responseExpiry = tokenResponse.res?.data?.expiry_date;
+    if (responseExpiry) {
+      if (typeof responseExpiry === "number") {
+        expiry = responseExpiry;
+      } else if (responseExpiry instanceof Date) {
+        expiry = responseExpiry.getTime();
+      } else if (typeof responseExpiry === "string") {
+        expiry = new Date(responseExpiry).getTime();
+      } else {
+        console.warn("Unknown expiry type", responseExpiry);
+      }
+    } else {
+      console.warn("No expiry date found in response", tokenResponse.res?.data);
+    }
+    return expiry;
+  }
+};
+
 // ai/anthropic/types.ts
 var AxAIAnthropicModel = /* @__PURE__ */ ((AxAIAnthropicModel2) => {
   AxAIAnthropicModel2["Claude35Sonnet"] = "claude-3-5-sonnet-latest";
@@ -800,6 +855,14 @@ var AxAIAnthropicModel = /* @__PURE__ */ ((AxAIAnthropicModel2) => {
   AxAIAnthropicModel2["ClaudeInstant12"] = "claude-instant-1.2";
   return AxAIAnthropicModel2;
 })(AxAIAnthropicModel || {});
+var AxAIAnthropicVertexModel = /* @__PURE__ */ ((AxAIAnthropicVertexModel2) => {
+  AxAIAnthropicVertexModel2["Claude35Haiku"] = "claude-3-5-haiku";
+  AxAIAnthropicVertexModel2["Claude35Sonnet"] = "claude-3-5-sonnet";
+  AxAIAnthropicVertexModel2["Claude35SonnetV2"] = "claude-3-5-sonnet-v2";
+  AxAIAnthropicVertexModel2["Claude3Haiku"] = "claude-3-haiku";
+  AxAIAnthropicVertexModel2["Claude3Opus"] = "claude-3-opus";
+  return AxAIAnthropicVertexModel2;
+})(AxAIAnthropicVertexModel || {});
 
 // ai/anthropic/info.ts
 var axModelInfoAnthropic = [
@@ -856,8 +919,9 @@ var axAIAnthropicDefaultConfig = () => structuredClone({
   ...axBaseAIDefaultConfig()
 });
 var AxAIAnthropicImpl = class {
-  constructor(config) {
+  constructor(config, isVertex) {
     this.config = config;
+    this.isVertex = isVertex;
   }
   getModelConfig() {
     const { config } = this;
@@ -876,9 +940,17 @@ var AxAIAnthropicImpl = class {
   }
   createChatReq = (req) => {
     const model = req.model;
-    const apiConfig = {
-      name: "/messages"
-    };
+    const stream = req.modelConfig?.stream ?? this.config.stream;
+    let apiConfig;
+    if (this.isVertex) {
+      apiConfig = {
+        name: stream ? `/models/${model}:streamRawPredict?alt=sse` : `/models/${model}:rawPredict`
+      };
+    } else {
+      apiConfig = {
+        name: "/messages"
+      };
+    }
     let toolsChoice;
     if (req.functionCall && req.functions && req.functions.length > 0) {
       if (typeof req.functionCall === "string") {
@@ -917,9 +989,8 @@ var AxAIAnthropicImpl = class {
         input_schema: v.parameters
       })
     );
-    const stream = req.modelConfig?.stream ?? this.config.stream;
     const reqValue = {
-      model,
+      ...this.isVertex ? { anthropic_version: "vertex-2023-10-16" } : { model },
       max_tokens: req.modelConfig?.maxTokens ?? this.config.maxTokens,
       stop_sequences: req.modelConfig?.stopSequences ?? this.config.stopSequences,
       temperature: req.modelConfig?.temperature ?? this.config.temperature,
@@ -1074,26 +1145,45 @@ var AxAIAnthropicImpl = class {
 var AxAIAnthropic = class extends AxBaseAI {
   constructor({
     apiKey,
+    projectId,
+    region,
     config,
     options,
     modelMap
   }) {
-    if (!apiKey || apiKey === "") {
-      throw new Error("Anthropic API key not set");
+    const isVertex = projectId !== void 0 && region !== void 0;
+    let apiURL;
+    let headers;
+    if (isVertex) {
+      apiURL = `https://${region}-aiplatform.googleapis.com/v1/projects/${projectId}/locations/${region}/publishers/anthropic/`;
+      if (apiKey) {
+        headers = async () => ({ Authorization: `Bearer ${apiKey}` });
+      } else {
+        const vertexAuth = new GoogleVertexAuth();
+        headers = async () => ({
+          Authorization: `Bearer ${await vertexAuth.getAccessToken()}`
+        });
+      }
+    } else {
+      if (!apiKey) {
+        throw new Error("Anthropic API key not set");
+      }
+      apiURL = "https://api.anthropic.com/v1";
+      headers = async () => ({
+        "anthropic-version": "2023-06-01",
+        "anthropic-beta": "prompt-caching-2024-07-31",
+        "x-api-key": apiKey
+      });
     }
     const _config = {
       ...axAIAnthropicDefaultConfig(),
       ...config
     };
-    const aiImpl = new AxAIAnthropicImpl(_config);
+    const aiImpl = new AxAIAnthropicImpl(_config, isVertex);
     super(aiImpl, {
       name: "Anthropic",
-      apiURL: "https://api.anthropic.com/v1",
-      headers: async () => ({
-        "anthropic-version": "2023-06-01",
-        "anthropic-beta": "prompt-caching-2024-07-31",
-        "x-api-key": apiKey
-      }),
+      apiURL,
+      headers,
       modelInfo: axModelInfoAnthropic,
       models: { model: _config.model },
       options,
@@ -2030,61 +2120,6 @@ var AxAIDeepSeek = class extends AxAIOpenAI {
   }
 };
 
-// ai/google-gemini/auth.ts
-import { GoogleAuth } from "google-auth-library";
-var GoogleVertexAuth = class {
-  auth;
-  client;
-  currentToken;
-  tokenExpiry;
-  constructor(config = {}) {
-    this.auth = new GoogleAuth({
-      scopes: ["https://www.googleapis.com/auth/cloud-platform"],
-      ...config
-    });
-  }
-  async getAuthenticatedClient() {
-    if (!this.client) {
-      this.client = await this.auth.getClient();
-    }
-    return this.client;
-  }
-  async getAccessToken() {
-    if (this.currentToken && this.tokenExpiry && Date.now() < this.tokenExpiry) {
-      return this.currentToken;
-    }
-    const client = await this.getAuthenticatedClient();
-    const tokenResponse = await client.getAccessToken();
-    this.currentToken = tokenResponse.token ?? void 0;
-    const expiry = this.getExpiry(tokenResponse);
-    const fiveMinutes = 5 * 60 * 1e3;
-    this.tokenExpiry = expiry - fiveMinutes;
-    return this.currentToken;
-  }
-  /**
-   * Get the expiry date from the token response.
-   */
-  getExpiry(tokenResponse) {
-    const oneHour = 3600 * 1e3;
-    let expiry = Date.now() + oneHour;
-    let responseExpiry = tokenResponse.res?.data?.expiry_date;
-    if (responseExpiry) {
-      if (typeof responseExpiry === "number") {
-        expiry = responseExpiry;
-      } else if (responseExpiry instanceof Date) {
-        expiry = responseExpiry.getTime();
-      } else if (typeof responseExpiry === "string") {
-        expiry = new Date(responseExpiry).getTime();
-      } else {
-        console.warn("Unknown expiry type", responseExpiry);
-      }
-    } else {
-      console.warn("No expiry date found in response", tokenResponse.res?.data);
-    }
-    return expiry;
-  }
-};
-
 // ai/google-gemini/types.ts
 var AxAIGoogleGeminiModel = /* @__PURE__ */ ((AxAIGoogleGeminiModel2) => {
   AxAIGoogleGeminiModel2["Gemini1Pro"] = "gemini-1.0-pro";
@@ -2174,11 +2209,10 @@ var axAIGoogleGeminiDefaultConfig = () => structuredClone({
   ...axBaseAIDefaultConfig()
 });
 var AxAIGoogleGeminiImpl = class {
-  constructor(config, isVertex, apiKey, keyFile, options) {
+  constructor(config, isVertex, apiKey, options) {
     this.config = config;
     this.isVertex = isVertex;
     this.apiKey = apiKey;
-    this.keyFile = keyFile;
     this.options = options;
   }
   getModelConfig() {
@@ -2455,7 +2489,6 @@ var AxAIGoogleGemini = class extends AxBaseAI {
     apiKey,
     projectId,
     region,
-    keyFile,
     config,
     options,
     modelMap
@@ -2468,9 +2501,7 @@ var AxAIGoogleGemini = class extends AxBaseAI {
       if (apiKey) {
         headers = async () => ({ Authorization: `Bearer ${apiKey}` });
       } else {
-        const vertexAuth = new GoogleVertexAuth({
-          keyFile
-        });
+        const vertexAuth = new GoogleVertexAuth();
         headers = async () => ({
           Authorization: `Bearer ${await vertexAuth.getAccessToken()}`
         });
@@ -2486,13 +2517,7 @@ var AxAIGoogleGemini = class extends AxBaseAI {
       ...axAIGoogleGeminiDefaultConfig(),
       ...config
     };
-    const aiImpl = new AxAIGoogleGeminiImpl(
-      _config,
-      isVertex,
-      apiKey,
-      keyFile,
-      options
-    );
+    const aiImpl = new AxAIGoogleGeminiImpl(_config, isVertex, apiKey, options);
     super(aiImpl, {
       name: "GoogleGeminiAI",
       apiURL,
@@ -4010,19 +4035,65 @@ var parseMarkdownList = (input) => {
   return list;
 };
 function mergeDeltas(base, delta) {
-  const merged = { ...base };
-  for (const key in delta) {
+  for (const key of Object.keys(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;
+    if ((baseValue === void 0 || Array.isArray(baseValue)) && Array.isArray(deltaValue)) {
+      base[key] = [...baseValue ?? [], ...deltaValue];
+    } else if ((baseValue === void 0 || typeof baseValue === "string") && typeof deltaValue === "string") {
+      base[key] = (baseValue ?? "") + deltaValue;
     } else {
-      merged[key] = deltaValue;
+      base[key] = deltaValue;
+    }
+  }
+  return base;
+}
+var LRUCache = class {
+  cache = /* @__PURE__ */ new Map();
+  maxSize;
+  constructor(maxSize) {
+    this.maxSize = maxSize;
+  }
+  get(key) {
+    const value = this.cache.get(key);
+    if (value) {
+      this.cache.delete(key);
+      this.cache.set(key, value);
+    }
+    return value;
+  }
+  set(key, value) {
+    if (this.cache.has(key)) {
+      this.cache.delete(key);
+    } else if (this.cache.size >= this.maxSize) {
+      const firstKey = this.cache.keys().next().value;
+      if (firstKey) {
+        this.cache.delete(firstKey);
+      }
+    }
+    this.cache.set(key, value);
+  }
+};
+var globalPrefixCache = new LRUCache(500);
+function matchesContent(content, prefix, startIndex = 0, prefixCache = globalPrefixCache) {
+  const exactMatchIndex = content.indexOf(prefix, startIndex);
+  if (exactMatchIndex !== -1) {
+    return exactMatchIndex;
+  }
+  const prefixes = prefixCache.get(prefix) ?? Array.from({ length: prefix.length }, (_, i) => prefix.slice(0, i + 1));
+  if (!prefixCache.get(prefix)) {
+    prefixCache.set(prefix, prefixes);
+  }
+  const contentEnd = content.slice(
+    Math.max(startIndex, content.length - prefix.length)
+  );
+  for (let i = 0; i < prefixes.length - 1; i++) {
+    const partialPrefix = prefixes[i];
+    if (contentEnd.endsWith(partialPrefix)) {
+      return -2;
     }
   }
-  return merged;
+  return -1;
 }
 
 // dsp/program.ts
@@ -4619,14 +4690,11 @@ function handleValidationError(mem, errorFields, ai, promptTemplate, sessionId)
   );
   mem.addTag("error");
   if (ai.getOptions().debug) {
-    process.stdout.write(
-      colorLog4.red(
-        `
+    const errors = errorFields.map((field) => `- ${field.title}: ${field.description}`).join("\n");
+    process.stdout.write(colorLog4.red(`
 Error Correction:
-${JSON.stringify(errorFields, null, 2)}
-`
-      )
-    );
+${errors}
+`));
   }
 }
 
@@ -4715,25 +4783,37 @@ var checkMissingRequiredFields = (xstate, values, currentIndex) => {
   }
 };
 var streamingExtractValues = (sig, values, xstate, content) => {
+  if (content.endsWith("\n")) {
+    return true;
+  }
   const fields = sig.getOutputFields();
   for (const [index, field] of fields.entries()) {
     if (field.name in values) {
       continue;
     }
     const prefix = field.title + ":";
-    const e = content.indexOf(prefix, xstate.s + 1);
-    if (e === -1) {
-      continue;
+    let e = matchesContent(content, prefix, xstate.s + 1);
+    switch (e) {
+      case -1:
+        continue;
+      // Field is not found, continue to the next field
+      case -2:
+        return true;
+    }
+    let prefixLen = prefix.length;
+    if (e - 1 >= 0 && content[e - 1] === "\n") {
+      e -= 1;
+      prefixLen += 1;
     }
     if (xstate.currField) {
-      const val = content.substring(xstate.s, e);
+      const val = content.substring(xstate.s, e).trim();
       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.s = e + prefixLen;
     xstate.currField = field;
     if (!xstate.extractedFields.includes(field)) {
       xstate.extractedFields.push(field);
@@ -4742,7 +4822,7 @@ var streamingExtractValues = (sig, values, xstate, content) => {
 };
 var streamingExtractFinalValue = (sig, values, xstate, content) => {
   if (xstate.currField) {
-    const val = content.substring(xstate.s);
+    const val = content.substring(xstate.s).trim();
     const parsedValue = validateAndParseFieldValue(xstate.currField, val);
     if (parsedValue !== void 0) {
       values[xstate.currField.name] = parsedValue;
@@ -4763,6 +4843,9 @@ var convertValueToType = (field, val) => {
       return v;
     }
     case "boolean": {
+      if (typeof val === "boolean") {
+        return val;
+      }
       const v = val.toLowerCase();
       if (v === "true") {
         return true;
@@ -4777,17 +4860,18 @@ var convertValueToType = (field, val) => {
     case "datetime":
       return parseLLMFriendlyDateTime(field, val);
     case "class":
-      if (field.type.classes && !field.type.classes.includes(val)) {
+      const className = val;
+      if (field.type.classes && !field.type.classes.includes(className)) {
         throw new Error(
           `Invalid class '${val}', expected one of the following: ${field.type.classes.join(", ")}`
         );
       }
-      return val;
+      return className;
     default:
       return val;
   }
 };
-function* streamingValues(sig, values, xstate, content) {
+function* streamValues(sig, values, xstate, content) {
   if (!xstate.currField) {
     return;
   }
@@ -4796,80 +4880,83 @@ function* streamingValues(sig, values, xstate, content) {
     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 pos = xstate.streamedIndex[fieldName] ?? 0;
+    const s = xstate.s + pos;
     const v = content.substring(s);
-    yield { [fieldName]: v };
-    xstate.streamedIndex[fieldName] = v.length;
+    yield { [fieldName]: pos === 0 ? v.trimStart() : v };
+    xstate.streamedIndex[fieldName] = pos + 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 s = xstate.streamedIndex[key] ?? 0;
       const v = value.slice(s);
       if (v) {
-        yield { [fieldName]: v };
-        xstate.streamedIndex[fieldName] = s + 1;
+        yield { [key]: v };
+        xstate.streamedIndex[key] = s + 1;
       }
       continue;
     }
-    if (!xstate.streamedIndex[fieldName]) {
-      yield { [fieldName]: value };
-      xstate.streamedIndex[fieldName] = 1;
+    if (!xstate.streamedIndex[key]) {
+      yield { [key]: value };
+      xstate.streamedIndex[key] = 1;
     }
   }
 }
 function validateAndParseFieldValue(field, fieldValue) {
-  const fv = fieldValue?.trim();
-  if (!fv || !fv || fv === "" || fv === "null" || fv === "NULL" || fv === "undefined") {
+  if (!fieldValue || fieldValue === "" || fieldValue === "null" || fieldValue === "NULL" || fieldValue === "undefined") {
     if (field.isOptional) {
       return;
     }
     throw new ValidationError({
       message: "Required field is missing",
       fields: [field],
-      value: fv
+      value: fieldValue
     });
   }
   let value;
   if (field.type?.name === "json") {
     try {
-      const text = extractBlock(fv);
+      const text = extractBlock(fieldValue);
       value = JSON5.parse(text);
       return value;
     } catch (e) {
       throw new ValidationError({
         message: "Invalid JSON: " + e.message,
         fields: [field],
-        value: fv
+        value: fieldValue
       });
     }
   }
   if (field.type?.isArray) {
     try {
       try {
-        value = JSON5.parse(fv);
+        value = JSON5.parse(fieldValue);
       } catch {
-        value = parseMarkdownList(fv);
+        value = parseMarkdownList(fieldValue);
       }
       if (!Array.isArray(value)) {
         throw new Error("Expected an array");
       }
     } catch (e) {
       throw new ValidationError({
-        message: "Invalid array: " + e.message,
+        message: "Invalid Array: " + e.message,
         fields: [field],
-        value: fv
+        value: fieldValue
       });
     }
   }
   try {
     if (Array.isArray(value)) {
       for (const [index, item] of value.entries()) {
-        value[index] = convertValueToType(field, item);
+        if (item !== void 0) {
+          const v = typeof item === "string" ? item.trim() : item;
+          value[index] = convertValueToType(field, v);
+        }
       }
     } else {
-      value = convertValueToType(field, fv);
+      value = convertValueToType(field, fieldValue);
     }
   } catch (e) {
     throw new ValidationError({
@@ -4878,7 +4965,10 @@ function validateAndParseFieldValue(field, fieldValue) {
       value: fieldValue
     });
   }
-  return value ?? fv;
+  if (typeof value === "string" && value === "") {
+    return void 0;
+  }
+  return value;
 }
 var extractBlock = (input) => {
   const jsonBlockPattern = /```([A-Za-z]+)?\s*([\s\S]*?)\s*```/g;
@@ -5160,7 +5250,10 @@ var AxGen = class extends AxProgramWithSignature {
   }) {
     const functionCalls = [];
     const values = {};
-    const xstate = { extractedFields: [], s: -1 };
+    const xstate = {
+      extractedFields: [],
+      s: -1
+    };
     let content = "";
     for await (const v of res) {
       const result = v.results[0];
@@ -5173,6 +5266,15 @@ var AxGen = class extends AxProgramWithSignature {
       if (result.content) {
         content += result.content;
         mem.updateResult({ name: result.name, content }, sessionId);
+        const skip = streamingExtractValues(
+          this.signature,
+          values,
+          xstate,
+          content
+        );
+        if (skip) {
+          continue;
+        }
         assertStreamingAssertions(
           this.streamingAsserts,
           values,
@@ -5180,9 +5282,8 @@ var AxGen = class extends AxProgramWithSignature {
           content,
           false
         );
-        streamingExtractValues(this.signature, values, xstate, content);
         assertAssertions(this.asserts, values);
-        yield* streamingValues(this.signature, values, xstate, content);
+        yield* streamValues(this.signature, values, xstate, content);
       }
       if (result.functionCalls) {
         mergeFunctionCalls(functionCalls, result.functionCalls);
@@ -5219,7 +5320,7 @@ var AxGen = class extends AxProgramWithSignature {
     );
     streamingExtractFinalValue(this.signature, values, xstate, content);
     assertAssertions(this.asserts, values);
-    return values;
+    yield* streamValues(this.signature, values, xstate, content);
   }
   async processResponse({
     ai,
@@ -5288,10 +5389,9 @@ var AxGen = class extends AxProgramWithSignature {
         try {
           const generator = this.forwardCore({ options, ai, mem });
           for await (const delta of generator) {
-            yield {
-              version: errCount,
-              delta
-            };
+            if (delta !== void 0) {
+              yield { version: errCount, delta };
+            }
           }
           const lastMemItem = mem.getLast(options?.sessionId);
           const shouldContinue = this.shouldContinueSteps(
@@ -5384,8 +5484,13 @@ var AxGen = class extends AxProgramWithSignature {
       ...options
     });
     let buffer = {};
-    for await (const delta of generator) {
-      buffer = mergeDeltas(buffer, delta.delta);
+    let currentVersion = 0;
+    for await (const item of generator) {
+      if (item.version !== currentVersion) {
+        buffer = {};
+      }
+      currentVersion = item.version;
+      buffer = mergeDeltas(buffer, item.delta);
     }
     this.trace = { ...values, ...buffer };
     return buffer;
@@ -7297,6 +7402,7 @@ export {
   AxAI,
   AxAIAnthropic,
   AxAIAnthropicModel,
+  AxAIAnthropicVertexModel,
   AxAIAzureOpenAI,
   AxAICohere,
   AxAICohereEmbedModel,