@ax-llm/ax 11.0.35 → 11.0.37

package/index.cjs

@@ -114,6 +114,7 @@ __export(index_exports, {
   AxStringUtil: () => AxStringUtil,
   AxTestPrompt: () => AxTestPrompt,
   axAIAnthropicDefaultConfig: () => axAIAnthropicDefaultConfig,
+  axAIAnthropicVertexDefaultConfig: () => axAIAnthropicVertexDefaultConfig,
   axAIAzureOpenAIBestConfig: () => axAIAzureOpenAIBestConfig,
   axAIAzureOpenAICreativeConfig: () => axAIAzureOpenAICreativeConfig,
   axAIAzureOpenAIDefaultConfig: () => axAIAzureOpenAIDefaultConfig,
@@ -885,7 +886,7 @@ var AxBaseAI = class {
     this.models = models;
     this.id = crypto.randomUUID();
     const model = this.getModel(defaults.model) ?? defaults.model;
-    const embedModel = this.getEmbedModel(defaults.embedModel);
+    const embedModel = this.getEmbedModel(defaults.embedModel) ?? defaults.embedModel;
     this.defaults = { model, embedModel };
     if (!defaults.model || typeof defaults.model !== "string" || defaults.model === "") {
       throw new Error("No model defined");
@@ -987,12 +988,6 @@ var AxBaseAI = class {
     }
     return models;
   }
-  getDefaultModels() {
-    return {
-      model: this.defaults.model,
-      embedModel: this.defaults.embedModel
-    };
-  }
   getName() {
     return this.name;
   }
@@ -1149,13 +1144,18 @@ var AxBaseAI = class {
       if (!this.aiImpl.createChatStreamResp) {
         throw new Error("generateChatResp not implemented");
       }
-      const respFn = this.aiImpl.createChatStreamResp;
+      const respFn = this.aiImpl.createChatStreamResp.bind(this);
       const wrappedRespFn = (state) => (resp) => {
         const res2 = respFn(resp, state);
         res2.sessionId = options?.sessionId;
-        if (res2.modelUsage) {
-          this.modelUsage = res2.modelUsage;
+        if (!res2.modelUsage) {
+          res2.modelUsage = {
+            ai: this.name,
+            model,
+            tokens: this.aiImpl.getTokenUsage()
+          };
         }
+        this.modelUsage = res2.modelUsage;
         if (span?.isRecording()) {
           setResponseAttr(res2, span);
         }
@@ -1261,17 +1261,19 @@ var AxBaseAI = class {
       );
       return res2;
     };
-    const resValue = this.rt ? await this.rt(fn, { embedModelUsage: this.embedModelUsage }) : await fn();
+    const resValue = this.rt ? await this.rt(fn, { modelUsage: this.embedModelUsage }) : await fn();
     const res = this.aiImpl.createEmbedResp(resValue);
     res.sessionId = options?.sessionId;
+    if (!res.modelUsage) {
+      res.modelUsage = {
+        ai: this.name,
+        model: embedModel,
+        tokens: this.aiImpl.getTokenUsage()
+      };
+    }
+    this.embedModelUsage = res.modelUsage;
     if (span?.isRecording()) {
-      if (res.modelUsage) {
-        this.embedModelUsage = res.modelUsage;
-        span.setAttributes({
-          [axSpanAttributes.LLM_USAGE_COMPLETION_TOKENS]: res.modelUsage.completionTokens ?? 0,
-          [axSpanAttributes.LLM_USAGE_PROMPT_TOKENS]: res.modelUsage.promptTokens
-        });
-      }
+      setResponseAttr(res, span);
     }
     span?.end();
     return res;
@@ -1298,8 +1300,8 @@ var AxBaseAI = class {
 function setResponseAttr(res, span) {
   if (res.modelUsage) {
     span.setAttributes({
-      [axSpanAttributes.LLM_USAGE_COMPLETION_TOKENS]: res.modelUsage.completionTokens ?? 0,
-      [axSpanAttributes.LLM_USAGE_PROMPT_TOKENS]: res.modelUsage.promptTokens
+      [axSpanAttributes.LLM_USAGE_COMPLETION_TOKENS]: res.modelUsage.tokens?.completionTokens ?? 0,
+      [axSpanAttributes.LLM_USAGE_PROMPT_TOKENS]: res.modelUsage.tokens?.promptTokens
     });
   }
 }
@@ -1354,14 +1356,14 @@ var AxAIAnthropicModel = /* @__PURE__ */ ((AxAIAnthropicModel2) => {
   AxAIAnthropicModel2["ClaudeInstant12"] = "claude-instant-1.2";
   return AxAIAnthropicModel2;
 })(AxAIAnthropicModel || {});
-var AxAIAnthropicVertexModel = /* @__PURE__ */ ((AxAIAnthropicVertexModel3) => {
-  AxAIAnthropicVertexModel3["Claude37Sonnet"] = "claude-3-7-sonnet";
-  AxAIAnthropicVertexModel3["Claude35Haiku"] = "claude-3-5-haiku";
-  AxAIAnthropicVertexModel3["Claude35Sonnet"] = "claude-3-5-sonnet";
-  AxAIAnthropicVertexModel3["Claude35SonnetV2"] = "claude-3-5-sonnet-v2";
-  AxAIAnthropicVertexModel3["Claude3Haiku"] = "claude-3-haiku";
-  AxAIAnthropicVertexModel3["Claude3Opus"] = "claude-3-opus";
-  return AxAIAnthropicVertexModel3;
+var AxAIAnthropicVertexModel = /* @__PURE__ */ ((AxAIAnthropicVertexModel2) => {
+  AxAIAnthropicVertexModel2["Claude37Sonnet"] = "claude-3-7-sonnet";
+  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
@@ -1415,7 +1417,11 @@ var axModelInfoAnthropic = [
 
 // ai/anthropic/api.ts
 var axAIAnthropicDefaultConfig = () => structuredClone({
-  model: "claude-3-5-sonnet-latest" /* Claude35Sonnet */,
+  model: "claude-3-7-sonnet-latest" /* Claude37Sonnet */,
+  ...axBaseAIDefaultConfig()
+});
+var axAIAnthropicVertexDefaultConfig = () => structuredClone({
+  model: "claude-3-7-sonnet" /* Claude37Sonnet */,
   ...axBaseAIDefaultConfig()
 });
 var AxAIAnthropicImpl = class {
@@ -1423,6 +1429,10 @@ var AxAIAnthropicImpl = class {
     this.config = config;
     this.isVertex = isVertex;
   }
+  tokensUsed;
+  getTokenUsage() {
+    return this.tokensUsed;
+  }
   getModelConfig() {
     const { config } = this;
     return {
@@ -1532,16 +1542,12 @@ var AxAIAnthropicImpl = class {
         finishReason
       };
     });
-    const modelUsage = {
+    this.tokensUsed = {
       promptTokens: resp.usage.input_tokens,
       completionTokens: resp.usage.output_tokens,
       totalTokens: resp.usage.input_tokens + resp.usage.output_tokens
     };
-    return {
-      results,
-      modelUsage,
-      remoteId: resp.id
-    };
+    return { results, remoteId: resp.id };
   };
   createChatStreamResp = (resp, state) => {
     if (!("type" in resp)) {
@@ -1558,15 +1564,12 @@ var AxAIAnthropicImpl = class {
     if (resp.type === "message_start") {
       const { message } = resp;
       const results = [{ content: "", id: message.id }];
-      const modelUsage = {
+      this.tokensUsed = {
         promptTokens: message.usage?.input_tokens ?? 0,
         completionTokens: message.usage?.output_tokens ?? 0,
         totalTokens: (message.usage?.input_tokens ?? 0) + (message.usage?.output_tokens ?? 0)
       };
-      return {
-        results,
-        modelUsage
-      };
+      return { results };
     }
     if (resp.type === "content_block_start") {
       const { content_block: contentBlock } = resp;
@@ -1623,19 +1626,15 @@ var AxAIAnthropicImpl = class {
     }
     if (resp.type === "message_delta") {
       const { delta, usage } = resp;
-      return {
-        results: [
-          {
-            content: "",
-            finishReason: mapFinishReason(delta.stop_reason)
-          }
-        ],
-        modelUsage: {
-          promptTokens: 0,
-          completionTokens: usage.output_tokens,
-          totalTokens: usage.output_tokens
-        }
+      this.tokensUsed = {
+        promptTokens: 0,
+        completionTokens: usage.output_tokens,
+        totalTokens: usage.output_tokens
       };
+      const results = [
+        { content: "", finishReason: mapFinishReason(delta.stop_reason) }
+      ];
+      return { results };
     }
     return {
       results: [{ content: "" }]
@@ -1826,6 +1825,7 @@ var AxAIOpenAIModel = /* @__PURE__ */ ((AxAIOpenAIModel2) => {
   AxAIOpenAIModel2["O4Mini"] = "o4-mini";
   AxAIOpenAIModel2["GPT4"] = "gpt-4";
   AxAIOpenAIModel2["GPT41"] = "gpt-4.1";
+  AxAIOpenAIModel2["GPT41Mini"] = "gpt-4.1-mini";
   AxAIOpenAIModel2["GPT4O"] = "gpt-4o";
   AxAIOpenAIModel2["GPT4OMini"] = "gpt-4o-mini";
   AxAIOpenAIModel2["GPT4ChatGPT4O"] = "chatgpt-4o-latest";
@@ -1882,6 +1882,12 @@ var axModelInfoOpenAI = [
     promptTokenCostPer1M: 2,
     completionTokenCostPer1M: 8
   },
+  {
+    name: "gpt-4.1-mini" /* GPT41Mini */,
+    currency: "usd",
+    promptTokenCostPer1M: 0.4,
+    completionTokenCostPer1M: 1.6
+  },
   {
     name: "gpt-4o" /* GPT4O */,
     currency: "usd",
@@ -1934,28 +1940,32 @@ var axModelInfoOpenAI = [
 
 // ai/openai/api.ts
 var axAIOpenAIDefaultConfig = () => structuredClone({
-  model: "gpt-4o" /* GPT4O */,
+  model: "gpt-4.1" /* GPT41 */,
   embedModel: "text-embedding-3-small" /* TextEmbedding3Small */,
   ...axBaseAIDefaultConfig()
 });
 var axAIOpenAIBestConfig = () => structuredClone({
   ...axAIOpenAIDefaultConfig(),
-  model: "gpt-4o" /* GPT4O */
+  model: "gpt-4.1" /* GPT41 */
 });
 var axAIOpenAICreativeConfig = () => structuredClone({
-  model: "gpt-4o" /* GPT4O */,
+  model: "gpt-4.1" /* GPT41 */,
   embedModel: "text-embedding-3-small" /* TextEmbedding3Small */,
   ...axBaseAIDefaultCreativeConfig()
 });
 var axAIOpenAIFastConfig = () => ({
   ...axAIOpenAIDefaultConfig(),
-  model: "gpt-4o-mini" /* GPT4OMini */
+  model: "gpt-4.1-mini" /* GPT41Mini */
 });
 var AxAIOpenAIImpl = class {
   constructor(config, streamingUsage) {
     this.config = config;
     this.streamingUsage = streamingUsage;
   }
+  tokensUsed;
+  getTokenUsage() {
+    return this.tokensUsed;
+  }
   getModelConfig() {
     const { config } = this;
     return {
@@ -2035,7 +2045,7 @@ var AxAIOpenAIImpl = class {
     if (error) {
       throw error;
     }
-    const modelUsage = usage ? {
+    this.tokensUsed = usage ? {
       promptTokens: usage.prompt_tokens,
       completionTokens: usage.completion_tokens,
       totalTokens: usage.total_tokens
@@ -2057,14 +2067,13 @@ var AxAIOpenAIImpl = class {
       };
     });
     return {
-      modelUsage,
       results,
       remoteId: id
     };
   }
   createChatStreamResp(resp, state) {
     const { id, usage, choices } = resp;
-    const modelUsage = usage ? {
+    this.tokensUsed = usage ? {
       promptTokens: usage.prompt_tokens,
       completionTokens: usage.completion_tokens,
       totalTokens: usage.total_tokens
@@ -2102,22 +2111,16 @@ var AxAIOpenAIImpl = class {
         };
       }
     );
-    return {
-      results,
-      modelUsage
-    };
+    return { results };
   }
   createEmbedResp(resp) {
     const { data, usage } = resp;
-    const modelUsage = usage ? {
+    this.tokensUsed = usage ? {
       promptTokens: usage.prompt_tokens,
       completionTokens: usage.completion_tokens,
       totalTokens: usage.total_tokens
     } : void 0;
-    return {
-      embeddings: data.map((v) => v.embedding),
-      modelUsage
-    };
+    return { embeddings: data.map((v) => v.embedding) };
   }
 };
 var mapFinishReason2 = (finishReason) => {
@@ -2382,6 +2385,10 @@ var AxAICohereImpl = class {
   constructor(config) {
     this.config = config;
   }
+  tokensUsed;
+  getTokenUsage() {
+    return this.tokensUsed;
+  }
   getModelConfig() {
     const { config } = this;
     return {
@@ -2473,7 +2480,7 @@ var AxAICohereImpl = class {
     return [apiConfig, reqValue];
   };
   createChatResp = (resp) => {
-    const modelUsage = resp.meta.billed_units ? {
+    this.tokensUsed = resp.meta.billed_units ? {
       promptTokens: resp.meta.billed_units.input_tokens,
       completionTokens: resp.meta.billed_units.output_tokens,
       totalTokens: resp.meta.billed_units.input_tokens + resp.meta.billed_units.output_tokens
@@ -2516,17 +2523,18 @@ var AxAICohereImpl = class {
         finishReason
       }
     ];
-    return {
-      results,
-      modelUsage,
-      remoteId: resp.response_id
-    };
+    return { results, remoteId: resp.response_id };
   };
   createChatStreamResp = (resp, state) => {
     const ss = state;
     if (resp.event_type === "stream-start") {
       ss.generation_id = resp.generation_id;
     }
+    this.tokensUsed = {
+      promptTokens: 0,
+      completionTokens: resp.meta.billed_units?.output_tokens ?? 0,
+      totalTokens: resp.meta.billed_units?.output_tokens ?? 0
+    };
     const { results } = this.createChatResp(resp);
     const result = results[0];
     if (!result) {
@@ -2820,13 +2828,13 @@ var safetySettings = [
 ];
 var axAIGoogleGeminiDefaultConfig = () => structuredClone({
   model: "gemini-2.0-flash" /* Gemini20Flash */,
-  embedModel: "text-embedding-004" /* TextEmbedding004 */,
+  embedModel: "text-embedding-005" /* TextEmbedding005 */,
   safetySettings,
   ...axBaseAIDefaultConfig()
 });
 var axAIGoogleGeminiDefaultCreativeConfig = () => structuredClone({
   model: "gemini-2.0-flash" /* Gemini20Flash */,
-  embedModel: "text-embedding-004" /* TextEmbedding004 */,
+  embedModel: "text-embedding-005" /* TextEmbedding005 */,
   safetySettings,
   ...axBaseAIDefaultCreativeConfig()
 });
@@ -2841,6 +2849,10 @@ var AxAIGoogleGeminiImpl = class {
       throw new Error("Auto truncate is not supported for GoogleGemini");
     }
   }
+  tokensUsed;
+  getTokenUsage() {
+    return this.tokensUsed;
+  }
   getModelConfig() {
     const { config } = this;
     return {
@@ -3103,18 +3115,14 @@ var AxAIGoogleGeminiImpl = class {
         return result;
       }
     );
-    let modelUsage;
     if (resp.usageMetadata) {
-      modelUsage = {
+      this.tokensUsed = {
         totalTokens: resp.usageMetadata.totalTokenCount,
         promptTokens: resp.usageMetadata.promptTokenCount,
         completionTokens: resp.usageMetadata.candidatesTokenCount
       };
     }
-    return {
-      results,
-      modelUsage
-    };
+    return { results };
   };
   createChatStreamResp = (resp) => {
     return this.createChatResp(resp);
@@ -3332,7 +3340,7 @@ var AxAIGroq = class extends AxAIOpenAIBase {
       debug: options?.debug
     });
     const rtFunc = async (func, info) => {
-      const totalTokens = info.modelUsage?.totalTokens || 0;
+      const totalTokens = info.modelUsage?.tokens?.totalTokens || 0;
       await rt.acquire(totalTokens);
       return await func();
     };
@@ -3362,6 +3370,10 @@ var AxAIHuggingFaceImpl = class {
   constructor(config) {
     this.config = config;
   }
+  tokensUsed;
+  getTokenUsage() {
+    return this.tokensUsed;
+  }
   getModelConfig() {
     const { config } = this;
     return {
@@ -3655,6 +3667,10 @@ var AxAIRekaImpl = class {
   constructor(config) {
     this.config = config;
   }
+  tokensUsed;
+  getTokenUsage() {
+    return this.tokensUsed;
+  }
   getModelConfig() {
     const { config } = this;
     return {
@@ -3695,7 +3711,7 @@ var AxAIRekaImpl = class {
   };
   createChatResp = (resp) => {
     const { id, usage, responses } = resp;
-    const modelUsage = usage ? {
+    this.tokensUsed = usage ? {
       promptTokens: usage.input_tokens,
       completionTokens: usage.output_tokens,
       totalTokens: usage.input_tokens + usage.output_tokens
@@ -3714,15 +3730,11 @@ var AxAIRekaImpl = class {
         finishReason
       };
     });
-    return {
-      modelUsage,
-      results,
-      remoteId: id
-    };
+    return { results, remoteId: id };
   };
   createChatStreamResp = (resp) => {
     const { id, usage, responses } = resp;
-    const modelUsage = usage ? {
+    this.tokensUsed = usage ? {
       promptTokens: usage.input_tokens,
       completionTokens: usage.output_tokens,
       totalTokens: usage.input_tokens + usage.output_tokens
@@ -3741,10 +3753,7 @@ var AxAIRekaImpl = class {
         finishReason
       };
     });
-    return {
-      results,
-      modelUsage
-    };
+    return { results };
   };
 };
 var mapFinishReason3 = (finishReason) => {
@@ -3930,9 +3939,6 @@ var AxAI = class {
   getModelList() {
     return this.ai.getModelList();
   }
-  getDefaultModels() {
-    return this.ai.getDefaultModels();
-  }
   getMetrics() {
     return this.ai.getMetrics();
   }
@@ -4869,9 +4875,15 @@ function mergeProgramUsage(usages) {
     }
     const currentUsage = usageMap[key];
     if (currentUsage) {
-      currentUsage.promptTokens += usage.promptTokens;
-      currentUsage.completionTokens += usage.completionTokens;
-      currentUsage.totalTokens += usage.totalTokens;
+      const tokens = currentUsage.tokens ?? {
+        promptTokens: 0,
+        completionTokens: 0,
+        totalTokens: 0
+      };
+      tokens.promptTokens += usage?.tokens?.promptTokens ?? 0;
+      tokens.completionTokens += usage?.tokens?.completionTokens ?? 0;
+      tokens.totalTokens += usage?.tokens?.totalTokens ?? 0;
+      currentUsage.tokens = tokens;
     }
   }
   return Object.values(usageMap);
@@ -6391,13 +6403,9 @@ var AxGen = class extends AxProgramWithSignature {
     mem,
     options
   }) {
-    const { sessionId, traceId, model, functions: _functions } = options ?? {};
+    const { sessionId, traceId, functions: _functions } = options ?? {};
     const fastFail = options?.fastFail ?? this.options?.fastFail;
-    const modelName = model ?? ai.getDefaultModels().model;
-    const usageInfo = {
-      ai: ai.getName(),
-      model: modelName
-    };
+    const model = options.model;
     const functions = _functions?.map((f) => "toFunction" in f ? f.toFunction() : f)?.flat();
     const res = await this.forwardSendRequest({
       ai,
@@ -6409,7 +6417,6 @@ var AxGen = class extends AxProgramWithSignature {
         ai,
         model,
         res,
-        usageInfo,
         mem,
         traceId,
         sessionId,
@@ -6421,7 +6428,6 @@ var AxGen = class extends AxProgramWithSignature {
         ai,
         model,
         res,
-        usageInfo,
         mem,
         traceId,
         sessionId,
@@ -6433,14 +6439,13 @@ var AxGen = class extends AxProgramWithSignature {
     ai,
     model,
     res,
-    usageInfo,
     mem,
     sessionId,
     traceId,
     functions,
     fastFail
   }) {
-    const streamingValidation = fastFail ?? ai.getFeatures().functionCot !== true;
+    const streamingValidation = fastFail ?? ai.getFeatures(model).functionCot !== true;
     const functionCalls = [];
     const values = {};
     const xstate = {
@@ -6455,7 +6460,7 @@ var AxGen = class extends AxProgramWithSignature {
         continue;
       }
       if (v.modelUsage) {
-        this.usage.push({ ...usageInfo, ...v.modelUsage });
+        this.usage.push(v.modelUsage);
       }
       if (result.functionCalls) {
         mergeFunctionCalls(functionCalls, result.functionCalls);
@@ -6556,7 +6561,6 @@ var AxGen = class extends AxProgramWithSignature {
   async processResponse({
     ai,
     res,
-    usageInfo,
     mem,
     sessionId,
     traceId,
@@ -6569,7 +6573,7 @@ var AxGen = class extends AxProgramWithSignature {
     }
     for (const result of results) {
       if (res.modelUsage) {
-        this.usage.push({ ...usageInfo, ...res.modelUsage });
+        this.usage.push(res.modelUsage);
       }
       mem.addResult(result, sessionId);
       if (result.functionCalls?.length) {
@@ -7161,9 +7165,6 @@ var AxBalancer = class _AxBalancer {
   getModelList() {
     return this.currentService.getModelList();
   }
-  getDefaultModels() {
-    return this.currentService.getDefaultModels();
-  }
   getNextService() {
     const cs = this.services[++this.currentServiceIndex];
     if (cs === void 0) {
@@ -9315,12 +9316,6 @@ var AxMockAIService = class {
   getModelList() {
     return this.config.models;
   }
-  getDefaultModels() {
-    return {
-      model: this.config.modelInfo?.name ?? "mock-model",
-      embedModel: this.config.embedModelInfo?.name
-    };
-  }
   getMetrics() {
     return this.metrics;
   }
@@ -9343,9 +9338,13 @@ var AxMockAIService = class {
         }
       ],
       modelUsage: {
-        promptTokens: 10,
-        completionTokens: 5,
-        totalTokens: 15
+        ai: this.getName(),
+        model: "mock-model",
+        tokens: {
+          promptTokens: 10,
+          completionTokens: 5,
+          totalTokens: 15
+        }
       }
     };
   }
@@ -9363,9 +9362,13 @@ var AxMockAIService = class {
     return this.config.embedResponse ?? {
       embeddings: [[0.1, 0.2, 0.3]],
       modelUsage: {
-        promptTokens: 5,
-        completionTokens: 0,
-        totalTokens: 5
+        ai: this.getName(),
+        model: "mock-model",
+        tokens: {
+          promptTokens: 5,
+          completionTokens: 0,
+          totalTokens: 5
+        }
       }
     };
   }
@@ -11329,11 +11332,6 @@ var AxMultiServiceRouter = class {
       }
     });
   }
-  getDefaultModels() {
-    throw new Error(
-      "getDefaultModels is not supported for multi-service router."
-    );
-  }
   /**
    * If a model key is provided, delegate to the corresponding service's features.
    * Otherwise, returns a default feature set.
@@ -11496,6 +11494,7 @@ var AxRAG = class extends AxChainOfThought {
   AxStringUtil,
   AxTestPrompt,
   axAIAnthropicDefaultConfig,
+  axAIAnthropicVertexDefaultConfig,
   axAIAzureOpenAIBestConfig,
   axAIAzureOpenAICreativeConfig,
   axAIAzureOpenAIDefaultConfig,