@ax-llm/ax 11.0.0 → 11.0.2

package/index.cjs

@@ -54,6 +54,7 @@ __export(index_exports, {
   AxAIMistralModel: () => AxAIMistralModel,
   AxAIOllama: () => AxAIOllama,
   AxAIOpenAI: () => AxAIOpenAI,
+  AxAIOpenAIBase: () => AxAIOpenAIBase,
   AxAIOpenAIEmbedModel: () => AxAIOpenAIEmbedModel,
   AxAIOpenAIModel: () => AxAIOpenAIModel,
   AxAIReka: () => AxAIReka,
@@ -750,8 +751,9 @@ var AxBaseAI = class {
     this.tracer = options.tracer;
     this.modelInfo = modelInfo;
     this.models = models;
+    this.id = crypto.randomUUID();
     const model = this.models?.find((v) => v.key === defaults.model)?.model ?? defaults.model;
-    const embedModel = this.models?.find((v) => v.key === defaults.embedModel)?.model ?? defaults.embedModel;
+    const embedModel = defaults.embedModel;
     this.defaults = { model, embedModel };
     if (!defaults.model || typeof defaults.model !== "string" || defaults.model === "") {
       throw new Error("No model defined");
@@ -769,6 +771,7 @@ var AxBaseAI = class {
   defaults;
   apiURL;
   name;
+  id;
   headers;
   supportFor;
   // Add private metrics tracking properties
@@ -803,6 +806,9 @@ var AxBaseAI = class {
   setName(name) {
     this.name = name;
   }
+  getId() {
+    return this.id;
+  }
   setAPIURL(apiURL) {
     this.apiURL = apiURL;
   }
@@ -848,8 +854,7 @@ var AxBaseAI = class {
     }
     const mi = getModelInfo({
       model: this.defaults.embedModel,
-      modelInfo: this.modelInfo,
-      models: this.models
+      modelInfo: this.modelInfo
     });
     return {
       ...mi,
@@ -1055,7 +1060,7 @@ var AxBaseAI = class {
     }
   }
   async _embed1(req, options) {
-    const embedModel = req.embedModel ? this.models?.find((v) => v.key === req.embedModel)?.model ?? req.embedModel : this.defaults.embedModel;
+    const embedModel = req.embedModel ?? this.defaults.embedModel;
     if (!embedModel) {
       throw new Error("No embed model defined");
     }
@@ -1066,7 +1071,7 @@ var AxBaseAI = class {
           kind: import_api2.SpanKind.SERVER,
           attributes: {
             [axSpanAttributes.LLM_SYSTEM]: this.name,
-            [axSpanAttributes.LLM_REQUEST_MODEL]: req.embedModel ?? this.defaults.embedModel
+            [axSpanAttributes.LLM_REQUEST_MODEL]: embedModel
           }
         },
         async (span) => {
@@ -1245,13 +1250,13 @@ 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;
+var AxAIAnthropicVertexModel = /* @__PURE__ */ ((AxAIAnthropicVertexModel3) => {
+  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;
 })(AxAIAnthropicVertexModel || {});
 
 // ai/anthropic/info.ts
@@ -2055,26 +2060,22 @@ function createMessages2(req) {
     }
   });
 }
-var AxAIOpenAI = class extends AxBaseAI {
+var AxAIOpenAIBase = class extends AxBaseAI {
   constructor({
     apiKey,
     config,
     options,
     apiURL,
-    modelInfo = axModelInfoOpenAI,
+    modelInfo,
     models
   }) {
     if (!apiKey || apiKey === "") {
       throw new Error("OpenAI API key not set");
     }
-    const _config = {
-      ...axAIOpenAIDefaultConfig(),
-      ...config
-    };
     const aiImpl = new AxAIOpenAIImpl(
-      _config,
+      config,
       options?.streamingUsage ?? true,
-      config?.dimensions
+      config.dimensions
     );
     super(aiImpl, {
       name: "OpenAI",
@@ -2082,8 +2083,8 @@ var AxAIOpenAI = class extends AxBaseAI {
       headers: async () => ({ Authorization: `Bearer ${apiKey}` }),
       modelInfo,
       defaults: {
-        model: _config.model,
-        embedModel: _config.embedModel
+        model: config.model,
+        embedModel: config.embedModel
       },
       options,
       supportFor: () => {
@@ -2096,10 +2097,33 @@ var AxAIOpenAI = class extends AxBaseAI {
 var isReasoningModel = (model) => ["o1-mini" /* O1Mini */, "o1" /* O1 */, "o3-mini" /* O3Mini */].includes(
   model
 );
+var AxAIOpenAI = class extends AxAIOpenAIBase {
+  constructor({
+    apiKey,
+    config,
+    options,
+    models
+  }) {
+    if (!apiKey || apiKey === "") {
+      throw new Error("OpenAI API key not set");
+    }
+    super({
+      apiKey,
+      config: {
+        ...axAIOpenAIDefaultConfig(),
+        ...config
+      },
+      options,
+      modelInfo: axModelInfoOpenAI,
+      models
+    });
+    super.setName("OpenAI");
+  }
+};
 
 // ai/azure-openai/api.ts
 var axAIAzureOpenAIDefaultConfig = axAIOpenAIDefaultConfig;
-var AxAIAzureOpenAI = class extends AxAIOpenAI {
+var AxAIAzureOpenAI = class extends AxAIOpenAIBase {
   constructor({
     apiKey,
     resourceName,
@@ -2122,7 +2146,13 @@ var AxAIAzureOpenAI = class extends AxAIOpenAI {
       ...axAIAzureOpenAIDefaultConfig(),
       ...config
     };
-    super({ apiKey, config: _config, options, models });
+    super({
+      apiKey,
+      config: _config,
+      options,
+      models,
+      modelInfo: axModelInfoOpenAI
+    });
     const host = resourceName.includes("://") ? resourceName : `https://${resourceName}.openai.azure.com/`;
     super.setName("Azure OpenAI");
     super.setAPIURL(
@@ -2485,7 +2515,7 @@ var axAIDeepSeekDefaultConfig = () => structuredClone({
   model: "deepseek-chat" /* DeepSeekChat */,
   ...axBaseAIDefaultConfig()
 });
-var AxAIDeepSeek = class extends AxAIOpenAI {
+var AxAIDeepSeek = class extends AxAIOpenAIBase {
   constructor({
     apiKey,
     config,
@@ -3049,7 +3079,7 @@ var axAIGroqDefaultConfig = () => structuredClone({
   model: "llama-3.3-70b-versatile" /* Llama33_70B */,
   ...axBaseAIDefaultConfig()
 });
-var AxAIGroq = class extends AxAIOpenAI {
+var AxAIGroq = class extends AxAIOpenAIBase {
   constructor({
     apiKey,
     config,
@@ -3294,7 +3324,7 @@ var axAIMistralDefaultConfig = () => structuredClone({
   model: "mistral-small-latest" /* MistralSmall */,
   ...axBaseAIDefaultConfig()
 });
-var AxAIMistral = class extends AxAIOpenAI {
+var AxAIMistral = class extends AxAIOpenAIBase {
   constructor({
     apiKey,
     config,
@@ -3326,7 +3356,7 @@ var axAIOllamaDefaultConfig = () => structuredClone({
   model: "nous-hermes2",
   embedModel: "all-minilm"
 });
-var AxAIOllama = class extends AxAIOpenAI {
+var AxAIOllama = class extends AxAIOpenAIBase {
   constructor({
     apiKey = "not-set",
     url = "http://localhost:11434/v1",
@@ -3343,7 +3373,8 @@ var AxAIOllama = class extends AxAIOpenAI {
       options,
       config: _config,
       apiURL: url,
-      models
+      models,
+      modelInfo: []
     });
     super.setName("Ollama");
   }
@@ -3580,7 +3611,7 @@ var axAITogetherDefaultConfig = () => structuredClone({
   model: "mistralai/Mixtral-8x7B-Instruct-v0.1",
   ...axBaseAIDefaultConfig()
 });
-var AxAITogether = class extends AxAIOpenAI {
+var AxAITogether = class extends AxAIOpenAIBase {
   constructor({
     apiKey,
     config,
@@ -3654,6 +3685,9 @@ var AxAI = class {
   getName() {
     return this.ai.getName();
   }
+  getId() {
+    return this.ai.getId();
+  }
   getModelInfo() {
     return this.ai.getModelInfo();
   }
@@ -6060,6 +6094,7 @@ var AxAgent = class {
   program;
   functions;
   agents;
+  disableSmartModelRouting;
   name;
   description;
   subAgentList;
@@ -6075,6 +6110,7 @@ var AxAgent = class {
     this.ai = ai;
     this.agents = agents;
     this.functions = functions;
+    this.disableSmartModelRouting = options?.disableSmartModelRouting;
     this.signature = new AxSignature(signature);
     this.signature.setDescription(description);
     if (!name || name.length < 5) {
@@ -6101,7 +6137,7 @@ var AxAgent = class {
       func: () => this.forward
     };
     const mm = ai?.getModelList();
-    if (mm) {
+    if (mm && !this.disableSmartModelRouting) {
       this.func.parameters = addModelParameter(this.func.parameters, mm);
     }
   }
@@ -6141,11 +6177,16 @@ var AxAgent = class {
       func: wrappedFunc
     };
   }
+  getFeatures() {
+    return {
+      canConfigureSmartModelRouting: this.ai !== void 0
+    };
+  }
   init(parentAi, options) {
     const ai = this.ai ?? parentAi;
     const mm = ai?.getModelList();
     const agentFuncs = this.agents?.map((a) => a.getFunction())?.map(
-      (f) => mm ? { ...f, parameters: addModelParameter(f.parameters, mm) } : f
+      (f) => mm && !this.disableSmartModelRouting && this.agents?.find((a) => a.getFunction().name === f.name)?.getFeatures().canConfigureSmartModelRouting ? { ...f, parameters: addModelParameter(f.parameters, mm) } : f
     );
     const functions = [
       ...options?.functions ?? this.functions ?? [],
@@ -6255,6 +6296,10 @@ var AxBalancer = class _AxBalancer {
   currentServiceIndex = 0;
   currentService;
   debug;
+  initialBackoffMs;
+  maxBackoffMs;
+  maxRetries;
+  serviceFailures = /* @__PURE__ */ new Map();
   constructor(services, options) {
     if (services.length === 0) {
       throw new Error("No AI services provided.");
@@ -6268,6 +6313,9 @@ var AxBalancer = class _AxBalancer {
     }
     this.currentService = cs;
     this.debug = options?.debug ?? true;
+    this.initialBackoffMs = options?.initialBackoffMs ?? 1e3;
+    this.maxBackoffMs = options?.maxBackoffMs ?? 32e3;
+    this.maxRetries = options?.maxRetries ?? 3;
   }
   /**
    * Service comparator that respects the input order of services.
@@ -6284,7 +6332,7 @@ var AxBalancer = class _AxBalancer {
     return aTotalCost - bTotalCost;
   };
   getModelList() {
-    throw new Error("Method not implemented.");
+    return this.currentService.getModelList();
   }
   getNextService() {
     const cs = this.services[++this.currentServiceIndex];
@@ -6305,6 +6353,9 @@ var AxBalancer = class _AxBalancer {
   getName() {
     return this.currentService.getName();
   }
+  getId() {
+    return this.currentService.getId();
+  }
   getModelInfo() {
     return this.currentService.getModelInfo();
   }
@@ -6317,11 +6368,56 @@ var AxBalancer = class _AxBalancer {
   getMetrics() {
     return this.currentService.getMetrics();
   }
+  canRetryService() {
+    const failure = this.serviceFailures.get(this.currentService.getId());
+    if (!failure) return true;
+    const { retries, lastFailureTime } = failure;
+    const timeSinceLastFailure = Date.now() - lastFailureTime;
+    const backoffMs = Math.min(
+      this.initialBackoffMs * Math.pow(2, retries),
+      this.maxBackoffMs
+    );
+    return timeSinceLastFailure >= backoffMs;
+  }
+  handleFailure() {
+    const failure = this.serviceFailures.get(this.currentService.getId());
+    const retries = (failure?.retries ?? 0) + 1;
+    this.serviceFailures.set(this.currentService.getId(), {
+      retries,
+      lastFailureTime: Date.now()
+    });
+    if (this.debug) {
+      console.warn(
+        `AxBalancer: Service ${this.currentService.getName()} failed (retry ${retries}/${this.maxRetries})`
+      );
+    }
+    if (retries >= this.maxRetries) {
+      const gotNextService = this.getNextService();
+      if (this.debug) {
+        console.warn(
+          `AxBalancer: Switching to service ${this.currentService.getName()}`
+        );
+      }
+      return gotNextService;
+    }
+    return true;
+  }
+  handleSuccess() {
+    this.serviceFailures.delete(this.currentService.getId());
+  }
   async chat(req, options) {
     this.reset();
     while (true) {
+      if (!this.canRetryService()) {
+        if (!this.getNextService()) {
+          throw new Error("All services exhausted");
+        }
+        continue;
+      }
       try {
-        return await this.currentService.chat(req, options);
+        const response = await this.currentService.chat(req, options);
+        this.handleSuccess();
+        return response;
       } catch (e) {
         if (!(e instanceof AxAIServiceError)) {
           throw e;
@@ -6342,38 +6438,29 @@ var AxBalancer = class _AxBalancer {
           default:
             throw e;
         }
-        if (this.debug) {
-          console.warn(
-            `AxBalancer: Service ${this.currentService.getName()} failed`,
-            e
-          );
-        }
-        if (!this.getNextService()) {
+        if (!this.handleFailure()) {
           throw e;
         }
-        if (this.debug) {
-          console.warn(
-            `AxBalancer: Switching to service ${this.currentService.getName()}`
-          );
-        }
       }
     }
   }
   async embed(req, options) {
     this.reset();
     while (true) {
+      if (!this.canRetryService()) {
+        if (!this.getNextService()) {
+          throw new Error("All services exhausted");
+        }
+        continue;
+      }
       try {
-        return await this.currentService.embed(req, options);
+        const response = await this.currentService.embed(req, options);
+        this.handleSuccess();
+        return response;
       } catch (e) {
-        if (this.debug) {
-          console.warn(`Service ${this.currentService.getName()} failed`);
-        }
-        if (!this.getNextService()) {
+        if (!this.handleFailure()) {
           throw e;
         }
-        if (this.debug) {
-          console.warn(`Switching to service ${this.currentService.getName()}`);
-        }
       }
     }
   }
@@ -7866,6 +7953,7 @@ var AxEmbeddingAdapter = class {
 var AxMockAIService = class {
   constructor(config = {}) {
     this.config = config;
+    this.config.id = this.config.id ?? crypto.randomUUID();
   }
   options = {};
   metrics = {
@@ -7881,6 +7969,9 @@ var AxMockAIService = class {
   getName() {
     return this.config.name ?? "mock-ai-service";
   }
+  getId() {
+    return this.config.id ?? "mock-ai-service-id";
+  }
   getModelInfo() {
     return {
       name: "mock-model",
@@ -8037,6 +8128,7 @@ var AxRAG = class extends AxChainOfThought {
   AxAIMistralModel,
   AxAIOllama,
   AxAIOpenAI,
+  AxAIOpenAIBase,
   AxAIOpenAIEmbedModel,
   AxAIOpenAIModel,
   AxAIReka,