@ax-llm/ax 10.0.37 → 10.0.39

package/index.cjs

@@ -649,7 +649,7 @@ var AxBaseAI = class {
         {
           name: apiConfig.name,
           url: this.apiURL,
-          headers: this.buildHeaders(apiConfig.headers),
+          headers: await this.buildHeaders(apiConfig.headers),
           stream: modelConfig.stream,
           debug: this.debug,
           fetch: this.fetch,
@@ -768,7 +768,7 @@ var AxBaseAI = class {
         {
           name: apiConfig.name,
           url: this.apiURL,
-          headers: this.buildHeaders(apiConfig.headers),
+          headers: await this.buildHeaders(apiConfig.headers),
           debug: this.debug,
           fetch: this.fetch,
           span
@@ -792,8 +792,8 @@ var AxBaseAI = class {
     span?.end();
     return res;
   }
-  buildHeaders(headers = {}) {
-    return { ...headers, ...this.headers };
+  async buildHeaders(headers = {}) {
+    return { ...headers, ...await this.headers() };
   }
 };
 var logChatRequest = (req) => {
@@ -1153,11 +1153,11 @@ var AxAIAnthropic = class extends AxBaseAI {
     super(aiImpl, {
       name: "Anthropic",
       apiURL: "https://api.anthropic.com/v1",
-      headers: {
+      headers: async () => ({
         "anthropic-version": "2023-06-01",
         "anthropic-beta": "prompt-caching-2024-07-31",
         "x-api-key": apiKey
-      },
+      }),
       modelInfo: axModelInfoAnthropic,
       models: { model: _config.model },
       options,
@@ -1662,7 +1662,7 @@ var AxAIOpenAI = class extends AxBaseAI {
     super(aiImpl, {
       name: "OpenAI",
       apiURL: apiURL ? apiURL : "https://api.openai.com/v1",
-      headers: { Authorization: `Bearer ${apiKey}` },
+      headers: async () => ({ Authorization: `Bearer ${apiKey}` }),
       modelInfo,
       models: {
         model: _config.model,
@@ -1714,7 +1714,7 @@ var AxAIAzureOpenAI = class extends AxAIOpenAI {
         host
       ).href
     );
-    super.setHeaders({ "api-key": apiKey });
+    super.setHeaders(async () => ({ "api-key": apiKey }));
   }
 };
 
@@ -1974,7 +1974,7 @@ var AxAICohere = class extends AxBaseAI {
     super(aiImpl, {
       name: "Cohere",
       apiURL: "https://api.cohere.ai/v1",
-      headers: { Authorization: `Bearer ${apiKey}` },
+      headers: async () => ({ Authorization: `Bearer ${apiKey}` }),
       modelInfo: axModelInfoCohere,
       models: { model: _config.model },
       supportFor: { functions: true, streaming: true },
@@ -2094,6 +2094,61 @@ var AxAIDeepSeek = class extends AxAIOpenAI {
   }
 };
 
+// ai/google-gemini/auth.ts
+var import_google_auth_library = require("google-auth-library");
+var GoogleVertexAuth = class {
+  auth;
+  client;
+  currentToken;
+  tokenExpiry;
+  constructor(config = {}) {
+    this.auth = new import_google_auth_library.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";
@@ -2106,6 +2161,7 @@ var AxAIGoogleGeminiModel = /* @__PURE__ */ ((AxAIGoogleGeminiModel2) => {
 })(AxAIGoogleGeminiModel || {});
 var AxAIGoogleGeminiEmbedModel = /* @__PURE__ */ ((AxAIGoogleGeminiEmbedModel2) => {
   AxAIGoogleGeminiEmbedModel2["TextEmbedding004"] = "text-embedding-004";
+  AxAIGoogleGeminiEmbedModel2["TextEmbedding005"] = "text-embedding-005";
   return AxAIGoogleGeminiEmbedModel2;
 })(AxAIGoogleGeminiEmbedModel || {});
 var AxAIGoogleGeminiSafetyCategory = /* @__PURE__ */ ((AxAIGoogleGeminiSafetyCategory2) => {
@@ -2182,10 +2238,11 @@ var axAIGoogleGeminiDefaultConfig = () => structuredClone({
   ...axBaseAIDefaultConfig()
 });
 var AxAIGoogleGeminiImpl = class {
-  constructor(config, apiKey, isVertex, options) {
+  constructor(config, isVertex, apiKey, keyFile, options) {
     this.config = config;
-    this.apiKey = apiKey;
     this.isVertex = isVertex;
+    this.apiKey = apiKey;
+    this.keyFile = keyFile;
     this.options = options;
   }
   getModelConfig() {
@@ -2212,7 +2269,7 @@ var AxAIGoogleGeminiImpl = class {
     const apiConfig = {
       name: stream ? `/models/${model}:streamGenerateContent?alt=sse` : `/models/${model}:generateContent`
     };
-    if (this.isVertex === false) {
+    if (!this.isVertex) {
       const pf = stream ? "&" : "?";
       apiConfig.name += `${pf}key=${this.apiKey}`;
     }
@@ -2359,15 +2416,28 @@ var AxAIGoogleGeminiImpl = class {
     if (!req.texts || req.texts.length === 0) {
       throw new Error("Embed texts is empty");
     }
-    const apiConfig = {
-      name: `/models/${model}:batchEmbedContents?key=${this.apiKey}`
-    };
-    const reqValue = {
-      requests: req.texts.map((text) => ({
-        model: "models/" + model,
-        content: { parts: [{ text }] }
-      }))
-    };
+    let apiConfig;
+    let reqValue;
+    if (this.isVertex) {
+      apiConfig = {
+        name: `/models/${model}:predict`
+      };
+      reqValue = {
+        instances: req.texts.map((text) => ({
+          content: text
+        }))
+      };
+    } else {
+      apiConfig = {
+        name: `/models/${model}:batchEmbedContents?key=${this.apiKey}`
+      };
+      reqValue = {
+        requests: req.texts.map((text) => ({
+          model: "models/" + model,
+          content: { parts: [{ text }] }
+        }))
+      };
+    }
     return [apiConfig, reqValue];
   };
   createChatResp = (resp) => {
@@ -2429,7 +2499,16 @@ var AxAIGoogleGeminiImpl = class {
     return this.createChatResp(resp);
   };
   createEmbedResp = (resp) => {
-    const embeddings = resp.embeddings.map((embedding) => embedding.values);
+    let embeddings;
+    if (this.isVertex) {
+      embeddings = resp.predictions.map(
+        (prediction) => prediction.embeddings.values
+      );
+    } else {
+      embeddings = resp.embeddings.map(
+        (embedding) => embedding.values
+      );
+    }
     return {
       embeddings
     };
@@ -2440,28 +2519,44 @@ var AxAIGoogleGemini = class extends AxBaseAI {
     apiKey,
     projectId,
     region,
+    keyFile,
     config,
     options,
     modelMap
   }) {
-    if (!apiKey || apiKey === "") {
-      throw new Error("GoogleGemini AI 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/google/`;
-      headers = { Authorization: `Bearer ${apiKey}` };
+      if (apiKey) {
+        headers = async () => ({ Authorization: `Bearer ${apiKey}` });
+      } else {
+        const vertexAuth = new GoogleVertexAuth({
+          keyFile
+        });
+        headers = async () => ({
+          Authorization: `Bearer ${await vertexAuth.getAccessToken()}`
+        });
+      }
     } else {
+      if (!apiKey) {
+        throw new Error("GoogleGemini AI API key not set");
+      }
       apiURL = "https://generativelanguage.googleapis.com/v1beta";
-      headers = {};
+      headers = async () => ({});
     }
     const _config = {
       ...axAIGoogleGeminiDefaultConfig(),
       ...config
     };
-    const aiImpl = new AxAIGoogleGeminiImpl(_config, apiKey, isVertex, options);
+    const aiImpl = new AxAIGoogleGeminiImpl(
+      _config,
+      isVertex,
+      apiKey,
+      keyFile,
+      options
+    );
     super(aiImpl, {
       name: "GoogleGeminiAI",
       apiURL,
@@ -2731,7 +2826,7 @@ var AxAIHuggingFace = class extends AxBaseAI {
     super(aiImpl, {
       name: "HuggingFace",
       apiURL: "https://api-inference.huggingface.co",
-      headers: { Authorization: `Bearer ${apiKey}` },
+      headers: async () => ({ Authorization: `Bearer ${apiKey}` }),
       modelInfo: axModelInfoHuggingFace,
       models: { model: _config.model },
       options,
@@ -3080,7 +3175,7 @@ var AxAIReka = class extends AxBaseAI {
     super(aiImpl, {
       name: "Reka",
       apiURL: apiURL ? apiURL : "https://api.reka.ai/v1/chat",
-      headers: { "X-Api-Key": apiKey },
+      headers: async () => ({ "X-Api-Key": apiKey }),
       modelInfo,
       models: {
         model: _config.model
@@ -5416,14 +5511,7 @@ var AxApacheTika = class {
 };
 
 // ai/balance.ts
-var axCostComparator = (a, b) => {
-  const aInfo = a.getModelInfo();
-  const bInfo = b.getModelInfo();
-  const aTotalCost = (aInfo.promptTokenCostPer1M || Infinity) + (aInfo.completionTokenCostPer1M || Infinity);
-  const bTotalCost = (bInfo.promptTokenCostPer1M || Infinity) + (bInfo.completionTokenCostPer1M || Infinity);
-  return aTotalCost - bTotalCost;
-};
-var AxBalancer = class {
+var AxBalancer = class _AxBalancer {
   services;
   currentServiceIndex = 0;
   currentService;
@@ -5431,13 +5519,29 @@ var AxBalancer = class {
     if (services.length === 0) {
       throw new Error("No AI services provided.");
     }
-    this.services = [...services].sort(options?.comparator ?? axCostComparator);
+    this.services = [...services].sort(
+      options?.comparator ?? _AxBalancer.costComparator
+    );
     const cs = this.services[this.currentServiceIndex];
     if (cs === void 0) {
       throw new Error("Error initializing the AI services.");
     }
     this.currentService = cs;
   }
+  /**
+   * Service comparator that respects the input order of services.
+   */
+  static inputOrderComparator = () => 0;
+  /**
+   * Service comparator that sorts services by cost.
+   */
+  static costComparator = (a, b) => {
+    const aInfo = a.getModelInfo();
+    const bInfo = b.getModelInfo();
+    const aTotalCost = (aInfo.promptTokenCostPer1M || Infinity) + (aInfo.completionTokenCostPer1M || Infinity);
+    const bTotalCost = (bInfo.promptTokenCostPer1M || Infinity) + (bInfo.completionTokenCostPer1M || Infinity);
+    return aTotalCost - bTotalCost;
+  };
   getModelMap() {
     throw new Error("Method not implemented.");
   }