@ax-llm/ax 11.0.36 → 11.0.38

package/index.js

@@ -735,7 +735,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");
@@ -754,6 +754,9 @@ var AxBaseAI = class {
   modelUsage;
   embedModelUsage;
   defaults;
+  lastUsedModelConfig;
+  lastUsedChatModel;
+  lastUsedEmbedModel;
   apiURL;
   name;
   id;
@@ -843,6 +846,15 @@ var AxBaseAI = class {
   getFeatures(model) {
     return typeof this.supportFor === "function" ? this.supportFor(model ?? this.defaults.model) : this.supportFor;
   }
+  getLastUsedChatModel() {
+    return this.lastUsedChatModel;
+  }
+  getLastUsedEmbedModel() {
+    return this.lastUsedEmbedModel;
+  }
+  getLastUsedModelConfig() {
+    return this.lastUsedModelConfig;
+  }
   // Method to calculate percentiles
   calculatePercentile(samples, percentile) {
     if (samples.length === 0) return 0;
@@ -965,6 +977,8 @@ var AxBaseAI = class {
       functions,
       modelConfig
     };
+    this.lastUsedChatModel = model;
+    this.lastUsedModelConfig = modelConfig;
     const fn = async () => {
       const [apiConfig, reqValue] = this.aiImpl.createChatReq(
         req,
@@ -1095,6 +1109,7 @@ var AxBaseAI = class {
       ...embedReq,
       embedModel
     };
+    this.lastUsedEmbedModel = embedModel;
     const fn = async () => {
       const [apiConfig, reqValue] = this.aiImpl.createEmbedReq(req);
       const res2 = await apiCall(
@@ -1122,13 +1137,7 @@ var AxBaseAI = class {
     }
     this.embedModelUsage = res.modelUsage;
     if (span?.isRecording()) {
-      if (res.modelUsage) {
-        this.embedModelUsage = res.modelUsage;
-        span.setAttributes({
-          [axSpanAttributes.LLM_USAGE_COMPLETION_TOKENS]: res.modelUsage.tokens?.completionTokens ?? 0,
-          [axSpanAttributes.LLM_USAGE_PROMPT_TOKENS]: res.modelUsage.tokens?.promptTokens
-        });
-      }
+      setResponseAttr(res, span);
     }
     span?.end();
     return res;
@@ -3794,6 +3803,15 @@ var AxAI = class {
   getModelList() {
     return this.ai.getModelList();
   }
+  getLastUsedChatModel() {
+    return this.ai.getLastUsedChatModel();
+  }
+  getLastUsedEmbedModel() {
+    return this.ai.getLastUsedEmbedModel();
+  }
+  getLastUsedModelConfig() {
+    return this.ai.getLastUsedModelConfig();
+  }
   getMetrics() {
     return this.ai.getMetrics();
   }
@@ -6365,7 +6383,10 @@ var AxGen = class extends AxProgramWithSignature {
         await assertAssertions(this.asserts, values);
       }
       if (result.finishReason === "length") {
-        throw new Error("Max tokens reached before completion");
+        throw new Error(
+          `Max tokens reached before completion
+Content: ${content}`
+        );
       }
     }
     const funcs = parseFunctionCalls(ai, functionCalls, values, model);
@@ -6460,7 +6481,10 @@ var AxGen = class extends AxProgramWithSignature {
         }
       }
       if (result.finishReason === "length") {
-        throw new Error("Max tokens reached before completion");
+        throw new Error(
+          `Max tokens reached before completion
+Content: ${result.content}`
+        );
       }
     }
     const publicValues = { ...values };
@@ -6525,7 +6549,7 @@ var AxGen = class extends AxProgramWithSignature {
             err = e;
           } else if (e instanceof AxAIServiceStreamTerminatedError) {
           } else {
-            throw e;
+            throw enhanceError(e, ai);
           }
           if (errorFields) {
             handleValidationError(
@@ -6540,7 +6564,7 @@ var AxGen = class extends AxProgramWithSignature {
       }
       throw new Error(`Unable to fix validation error: ${err?.toString()}`);
     }
-    throw new Error(`Max steps reached: ${maxSteps}`);
+    throw enhanceError(new Error(`Max steps reached: ${maxSteps}`), ai);
   }
   shouldContinueSteps(lastMemItem, stopFunction) {
     const stopFunctionExecuted = stopFunction && this.functionsExecuted.has(stopFunction);
@@ -6611,6 +6635,20 @@ var AxGen = class extends AxProgramWithSignature {
     });
   }
 };
+function enhanceError(e, ai) {
+  const originalError = e instanceof Error ? e : new Error(String(e));
+  const model = ai.getLastUsedChatModel();
+  const modelConfig = ai.getLastUsedModelConfig();
+  const details = [
+    `model=${model}`,
+    `maxTokens=${modelConfig?.maxTokens ?? "N/A"}`,
+    `streaming=${modelConfig?.stream ?? false}`
+  ].join(", ");
+  const enhancedError = new Error();
+  enhancedError.stack = `Generate Failed: ${details}
+${originalError.stack}`;
+  return enhancedError;
+}
 
 // prompts/agent.ts
 function processChildAgentFunction(childFunction, parentValues, parentInputKeys, modelList, options) {
@@ -6991,6 +7029,15 @@ var AxBalancer = class _AxBalancer {
     this.maxBackoffMs = options?.maxBackoffMs ?? 32e3;
     this.maxRetries = options?.maxRetries ?? 3;
   }
+  getLastUsedChatModel() {
+    return this.currentService.getLastUsedChatModel();
+  }
+  getLastUsedEmbedModel() {
+    return this.currentService.getLastUsedEmbedModel();
+  }
+  getLastUsedModelConfig() {
+    return this.currentService.getLastUsedModelConfig();
+  }
   /**
    * Service comparator that respects the input order of services.
    */
@@ -9155,6 +9202,15 @@ var AxMockAIService = class {
       embed: { count: 0, rate: 0, total: 0 }
     }
   };
+  getLastUsedChatModel() {
+    throw new Error("Method not implemented.");
+  }
+  getLastUsedEmbedModel() {
+    throw new Error("Method not implemented.");
+  }
+  getLastUsedModelConfig() {
+    throw new Error("Method not implemented.");
+  }
   getName() {
     return this.config.name ?? "mock-ai-service";
   }
@@ -11064,6 +11120,7 @@ var AxMCPStdioTransport = class {
 // ai/multiservice.ts
 var AxMultiServiceRouter = class {
   options;
+  lastUsedService;
   services = /* @__PURE__ */ new Map();
   /**
    * Constructs a new multi-service router.
@@ -11122,6 +11179,15 @@ var AxMultiServiceRouter = class {
       }
     }
   }
+  getLastUsedChatModel() {
+    return this.lastUsedService?.getLastUsedChatModel();
+  }
+  getLastUsedEmbedModel() {
+    return this.lastUsedService?.getLastUsedEmbedModel();
+  }
+  getLastUsedModelConfig() {
+    return this.lastUsedService?.getLastUsedModelConfig();
+  }
   /**
    * Delegates the chat call to the service matching the provided model key.
    */
@@ -11134,6 +11200,7 @@ var AxMultiServiceRouter = class {
     if (!item) {
       throw new Error(`No service found for model key: ${modelKey}`);
     }
+    this.lastUsedService = item.service;
     if (!item.model) {
       const { model, ...reqWithoutModel } = req;
       return await item.service.chat(reqWithoutModel, options);
@@ -11152,6 +11219,7 @@ var AxMultiServiceRouter = class {
     if (!item) {
       throw new Error(`No service found for embed model key: ${embedModelKey}`);
     }
+    this.lastUsedService = item.service;
     if (!item.model) {
       const { embedModel, ...reqWithoutEmbedModel } = req;
       return await item.service.embed(reqWithoutEmbedModel, options);
@@ -11206,11 +11274,17 @@ var AxMultiServiceRouter = class {
    * or falls back to the first service if none has been used.
    */
   getMetrics() {
-    const service = this.services.values().next().value;
-    if (!service) {
+    let serviceInstance = this.lastUsedService;
+    if (!serviceInstance) {
+      const firstServiceEntry = this.services.values().next().value;
+      if (firstServiceEntry) {
+        serviceInstance = "service" in firstServiceEntry ? firstServiceEntry.service : firstServiceEntry;
+      }
+    }
+    if (!serviceInstance) {
       throw new Error("No service available to get metrics.");
     }
-    return service.service.getMetrics();
+    return serviceInstance.getMetrics();
   }
   /**
    * Sets options on all underlying services.