@ax-llm/ax 11.0.20 → 11.0.22

package/index.js

@@ -1,26 +1,6 @@
 // ai/base.ts
 import { SpanKind } from "@opentelemetry/api";
 
-// dsp/modelinfo.ts
-function getModelInfo({
-  model,
-  modelInfo,
-  models
-}) {
-  const mappedModel = models?.find((v) => v.key === model)?.model ?? model;
-  const exactMatch = modelInfo.find((v) => v.name === model);
-  if (exactMatch) return exactMatch;
-  const normalizedName = mappedModel.replace(/^(anthropic\.|openai\.)/, "").replace(/-latest$/, "").replace(/-\d{8}$/, "").replace(/-v\d+:\d+$/, "").replace(/@\d{8}$/, "").replace(/-\d{2,}(-[a-zA-Z0-9-]+)?$/, "").replace(/-v\d+@\d{8}$/, "").replace(/-v\d+$/, "");
-  const normalizedMatch = modelInfo.find((v) => v.name === normalizedName);
-  if (normalizedMatch) return normalizedMatch;
-  return {
-    name: model,
-    currency: "usd",
-    promptTokenCostPer1M: 0,
-    completionTokenCostPer1M: 0
-  };
-}
-
 // trace/trace.ts
 var axSpanAttributes = {
   // LLM
@@ -746,6 +726,9 @@ var AxBaseAI = class {
       throw new Error("No model defined");
     }
     this.setOptions(options);
+    if (models) {
+      validateModels(models);
+    }
   }
   debug = false;
   rt;
@@ -824,33 +807,19 @@ var AxBaseAI = class {
       tracer: this.tracer
     };
   }
-  getModelInfo() {
-    const mi = getModelInfo({
-      model: this.defaults.model,
-      modelInfo: this.modelInfo,
-      models: this.models
-    });
-    return {
-      ...mi,
-      provider: this.name
-    };
+  getModelList() {
+    return this.models?.filter((model) => !model.isInternal)?.map((model) => ({
+      key: model.key,
+      description: model.description,
+      model: model.model
+    }));
   }
-  getEmbedModelInfo() {
-    if (!this.defaults.embedModel) {
-      return;
-    }
-    const mi = getModelInfo({
-      model: this.defaults.embedModel,
-      modelInfo: this.modelInfo
-    });
+  getDefaultModels() {
     return {
-      ...mi,
-      provider: this.name
+      model: this.defaults.model,
+      embedModel: this.defaults.embedModel
     };
   }
-  getModelList() {
-    return this.models;
-  }
   getName() {
     return this.name;
   }
@@ -1136,14 +1105,25 @@ var AxBaseAI = class {
     return { ...headers, ...await this.headers() };
   }
 };
-var setResponseAttr = (res, span) => {
+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
     });
   }
-};
+}
+function validateModels(models) {
+  const keys = /* @__PURE__ */ new Set();
+  for (const model of models) {
+    if (keys.has(model.key)) {
+      throw new Error(
+        `Duplicate model key detected: "${model.key}". Each model key must be unique.`
+      );
+    }
+    keys.add(model.key);
+  }
+}
 
 // ai/google-vertex/auth.ts
 import { GoogleAuth } from "google-auth-library";
@@ -1799,7 +1779,7 @@ var AxAIOpenAIImpl = class {
       response_format: this.config?.responseFormat ? { type: this.config?.responseFormat } : void 0,
       tools,
       tool_choice: toolsChoice,
-      max_tokens: req.modelConfig?.maxTokens ?? this.config.maxTokens ?? 500,
+      max_completion_tokens: req.modelConfig?.maxTokens ?? this.config.maxTokens ?? 500,
       temperature: req.modelConfig?.temperature ?? this.config.temperature,
       top_p: req.modelConfig?.topP ?? this.config.topP ?? 1,
       n: req.modelConfig?.n ?? this.config.n,
@@ -2596,10 +2576,10 @@ var axAIGoogleGeminiDefaultConfig = () => structuredClone({
   ...axBaseAIDefaultConfig()
 });
 var AxAIGoogleGeminiImpl = class {
-  constructor(config, isVertex, endpoint, apiKey, options) {
+  constructor(config, isVertex, endpointId, apiKey, options) {
     this.config = config;
     this.isVertex = isVertex;
-    this.endpoint = endpoint;
+    this.endpointId = endpointId;
     this.apiKey = apiKey;
     this.options = options;
   }
@@ -2625,9 +2605,9 @@ var AxAIGoogleGeminiImpl = class {
       throw new Error("Chat prompt is empty");
     }
     let apiConfig;
-    if (this.endpoint) {
+    if (this.endpointId) {
       apiConfig = {
-        name: stream ? `/${this.endpoint}:streamGenerateContent?alt=sse` : `/${this.endpoint}:generateContent`
+        name: stream ? `/${this.endpointId}:streamGenerateContent?alt=sse` : `/${this.endpointId}:generateContent`
       };
     } else {
       apiConfig = {
@@ -2786,9 +2766,9 @@ var AxAIGoogleGeminiImpl = class {
     let apiConfig;
     let reqValue;
     if (this.isVertex) {
-      if (this.endpoint) {
+      if (this.endpointId) {
         apiConfig = {
-          name: `/${this.endpoint}:predict`
+          name: `/${this.endpointId}:predict`
         };
       } else {
         apiConfig = {
@@ -2892,7 +2872,7 @@ var AxAIGoogleGemini = class extends AxBaseAI {
     apiKey,
     projectId,
     region,
-    endpoint,
+    endpointId,
     config,
     options,
     models
@@ -2902,7 +2882,7 @@ var AxAIGoogleGemini = class extends AxBaseAI {
     let headers;
     if (isVertex) {
       let path;
-      if (endpoint) {
+      if (endpointId) {
         path = "endpoints";
       } else {
         path = "publishers/google";
@@ -2930,7 +2910,7 @@ var AxAIGoogleGemini = class extends AxBaseAI {
     const aiImpl = new AxAIGoogleGeminiImpl(
       _config,
       isVertex,
-      endpoint,
+      endpointId,
       apiKey,
       options
     );
@@ -3651,18 +3631,15 @@ var AxAI = class {
   getId() {
     return this.ai.getId();
   }
-  getModelInfo() {
-    return this.ai.getModelInfo();
-  }
-  getEmbedModelInfo() {
-    return this.ai.getEmbedModelInfo();
-  }
   getFeatures(model) {
     return this.ai.getFeatures(model);
   }
   getModelList() {
     return this.ai.getModelList();
   }
+  getDefaultModels() {
+    return this.ai.getDefaultModels();
+  }
   getMetrics() {
     return this.ai.getMetrics();
   }
@@ -3899,10 +3876,11 @@ var AxAssertionError = class extends Error {
   }
   getFixingInstructions = () => {
     const extraFields = [];
+    const message = this.message.trim();
     extraFields.push({
       name: "error",
-      title: "Error In Output",
-      description: `You must follow the following instructions, "${this.message}".`
+      title: "Follow these instructions",
+      description: message + (message.endsWith(".") ? "" : ".")
     });
     return extraFields;
   };
@@ -5639,7 +5617,8 @@ async function processFieldProcessors(fieldProcessors, values, mem, sessionId) {
     if (values[processor.field.name] === void 0) {
       continue;
     }
-    const result = await processor.process(values[processor.field.name], {
+    const processFn = processor.process;
+    const result = await processFn(values[processor.field.name], {
       sessionId,
       values,
       done: true
@@ -5657,7 +5636,8 @@ async function processStreamingFieldProcessors(fieldProcessors, content, xstate,
       value = value.replace(/^[ ]*```[a-zA-Z0-9]*\n\s*/, "");
       value = value.replace(/\s*```\s*$/, "");
     }
-    const result = await processor.process(value, {
+    const processFn = processor.process;
+    const result = await processFn(value, {
       sessionId,
       values,
       done
@@ -5927,7 +5907,7 @@ var AxGen = class extends AxProgramWithSignature {
   addStreamingAssert = (fieldName, fn, message) => {
     this.streamingAsserts.push({ fieldName, fn, message });
   };
-  addFieldProcessor = (fieldName, fn, streaming = false) => {
+  addFieldProcessorInternal = (fieldName, fn, streaming = false) => {
     const field = this.signature.getOutputFields().find((f) => f.name === fieldName);
     if (!field) {
       throw new Error(`addFieldProcessor: field ${fieldName} not found`);
@@ -5945,6 +5925,12 @@ var AxGen = class extends AxProgramWithSignature {
       this.fieldProcessors.push({ field, process: fn });
     }
   };
+  addStreamingFieldProcessor = (fieldName, fn) => {
+    this.addFieldProcessorInternal(fieldName, fn, true);
+  };
+  addFieldProcessor = (fieldName, fn) => {
+    this.addFieldProcessorInternal(fieldName, fn, false);
+  };
   async forwardSendRequest({
     ai,
     mem,
@@ -5990,9 +5976,10 @@ var AxGen = class extends AxProgramWithSignature {
   }) {
     const { sessionId, traceId, model, functions } = options ?? {};
     const fastFail = options?.fastFail ?? this.options?.fastFail;
+    const modelName = model ?? ai.getDefaultModels().model;
     const usageInfo = {
       ai: ai.getName(),
-      model: ai.getModelInfo().name
+      model: modelName
     };
     const res = await this.forwardSendRequest({
       ai,
@@ -6076,7 +6063,9 @@ var AxGen = class extends AxProgramWithSignature {
           content,
           streamingValidation
         );
-        assertStreamingAssertions(this.streamingAsserts, xstate, content);
+        if (this.streamingAsserts.length !== 0) {
+          assertStreamingAssertions(this.streamingAsserts, xstate, content);
+        }
         if (this.streamingFieldProcessors.length !== 0) {
           await processStreamingFieldProcessors(
             this.streamingFieldProcessors,
@@ -6702,8 +6691,9 @@ var AxBalancer = class _AxBalancer {
     if (services.length === 0) {
       throw new Error("No AI services provided.");
     }
+    validateModels2(services);
     this.services = [...services].sort(
-      options?.comparator ?? _AxBalancer.costComparator
+      options?.comparator ?? _AxBalancer.metricComparator
     );
     const cs = this.services[this.currentServiceIndex];
     if (cs === void 0) {
@@ -6722,16 +6712,31 @@ var AxBalancer = class _AxBalancer {
   /**
    * 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;
+  // Requires a rethink
+  /*
+    public static costComparator = (a: AxAIService, b: AxAIService) => {
+      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
+    }
+    */
+  static metricComparator = (a, b) => {
+    const aMetrics = a.getMetrics();
+    const bMetrics = b.getMetrics();
+    return aMetrics.latency.chat.mean - bMetrics.latency.chat.mean;
   };
   getModelList() {
     return this.currentService.getModelList();
   }
+  getDefaultModels() {
+    return this.currentService.getDefaultModels();
+  }
   getNextService() {
     const cs = this.services[++this.currentServiceIndex];
     if (cs === void 0) {
@@ -6754,12 +6759,6 @@ var AxBalancer = class _AxBalancer {
   getId() {
     return this.currentService.getId();
   }
-  getModelInfo() {
-    return this.currentService.getModelInfo();
-  }
-  getEmbedModelInfo() {
-    return this.currentService.getEmbedModelInfo();
-  }
   getFeatures(model) {
     return this.currentService.getFeatures(model);
   }
@@ -6869,6 +6868,46 @@ var AxBalancer = class _AxBalancer {
     return this.currentService.getOptions();
   }
 };
+function validateModels2(services) {
+  const serviceWithModel = services.find(
+    (service) => service.getModelList() !== void 0
+  );
+  if (!serviceWithModel) {
+    return;
+  }
+  const referenceModelList = serviceWithModel.getModelList();
+  if (!referenceModelList) {
+    throw new Error("No model list found in any service.");
+  }
+  const referenceKeys = new Set(referenceModelList.map((model) => model.key));
+  for (let i = 0; i < services.length; i++) {
+    const service = services[i];
+    if (!service) {
+      throw new Error(`Service at index ${i} is undefined`);
+    }
+    const modelList = service.getModelList();
+    if (!modelList) {
+      throw new Error(
+        `Service at index ${i} (${service.getName()}) has no model list while another service does.`
+      );
+    }
+    const serviceKeys = new Set(modelList.map((model) => model.key));
+    for (const key of referenceKeys) {
+      if (!serviceKeys.has(key)) {
+        throw new Error(
+          `Service at index ${i} (${service.getName()}) is missing model "${key}"`
+        );
+      }
+    }
+    for (const key of serviceKeys) {
+      if (!referenceKeys.has(key)) {
+        throw new Error(
+          `Service at index ${i} (${service.getName()}) has extra model "${key}"`
+        );
+      }
+    }
+  }
+}
 
 // dsp/optimize.ts
 var AxBootstrapFewShot = class {
@@ -8100,18 +8139,6 @@ var AxMockAIService = class {
   getId() {
     return this.config.id ?? "mock-ai-service-id";
   }
-  getModelInfo() {
-    return {
-      name: "mock-model",
-      provider: "mock-provider",
-      promptTokenCostPer1M: 100,
-      completionTokenCostPer1M: 100,
-      ...this.config.modelInfo
-    };
-  }
-  getEmbedModelInfo() {
-    return this.config.embedModelInfo;
-  }
   // eslint-disable-next-line @typescript-eslint/no-unused-vars
   getFeatures(_model) {
     return {
@@ -8122,6 +8149,12 @@ 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;
   }
@@ -8199,7 +8232,7 @@ var AxMockAIService = class {
 
 // dsp/router.ts
 var colorLog6 = new ColorLog();
-var AxRoute = class {
+var AxSimpleClassifierClass = class {
   name;
   context;
   constructor(name, context) {
@@ -8213,7 +8246,7 @@ var AxRoute = class {
     return this.context;
   }
 };
-var AxRouter = class {
+var AxSimpleClassifier = class {
   ai;
   db;
   debug;
@@ -8227,12 +8260,12 @@ var AxRouter = class {
   setState(state) {
     this.db.setDB(state);
   }
-  setRoutes = async (routes) => {
-    for (const ro of routes) {
-      const ret = await this.ai.embed({ texts: ro.getContext() });
+  setClasses = async (classes) => {
+    for (const c of classes) {
+      const ret = await this.ai.embed({ texts: c.getContext() });
       await this.db.upsert({
-        id: ro.getName(),
-        table: "routes",
+        id: c.getName(),
+        table: "classes",
         values: ret.embeddings[0]
       });
     }
@@ -8240,7 +8273,7 @@ var AxRouter = class {
   async forward(text, options) {
     const { embeddings } = await this.ai.embed({ texts: [text] });
     const matches = await this.db.query({
-      table: "routes",
+      table: "classes",
       values: embeddings[0]
     });
     let m = matches.matches;
@@ -8255,11 +8288,11 @@ var AxRouter = class {
         )
       );
     }
-    const route = m.at(0);
-    if (!route) {
+    const matchedClass = m.at(0);
+    if (!matchedClass) {
       return "";
     }
-    return route.id;
+    return matchedClass.id;
   }
   setOptions(options) {
     if (typeof options.debug === "boolean") {
@@ -8466,6 +8499,160 @@ var AxEmbeddingAdapter = class {
   }
 };
 
+// ai/multiservice.ts
+var AxMultiServiceRouter = class {
+  services = /* @__PURE__ */ new Map();
+  /**
+   * Constructs a new multi-service router.
+   * It validates that each service provides a unique set of model keys,
+   * then builds a lookup (map) for routing the chat/embed requests.
+   */
+  constructor(services) {
+    if (services.length === 0) {
+      throw new Error("No AI services provided.");
+    }
+    for (const [index, item] of services.entries()) {
+      const isKeyBased = "key" in item;
+      if (isKeyBased) {
+        if (this.services.has(item.key)) {
+          throw new Error(`Duplicate model key: ${item.key}`);
+        }
+        const { service, description, isInternal } = item;
+        this.services.set(item.key, {
+          service,
+          description,
+          isInternal,
+          model: item.service.getDefaultModels().model,
+          useDefaultModel: true
+        });
+      } else {
+        const modelList = item.getModelList();
+        if (!modelList) {
+          throw new Error(
+            `Service ${index} \`${item.getName()}\` has no model list.`
+          );
+        }
+        for (const { key, description, model } of modelList ?? []) {
+          if (this.services.has(key)) {
+            const otherService = this.services.get(key)?.service;
+            throw new Error(
+              `Service ${index} \`${item.getName()}\` has duplicate model key: ${key} as service ${otherService?.getName()}`
+            );
+          }
+          this.services.set(key, {
+            description,
+            service: item,
+            model
+          });
+        }
+      }
+    }
+  }
+  /**
+   * Delegates the chat call to the service matching the provided model key.
+   */
+  async chat(req, options) {
+    const modelKey = req.model;
+    if (!modelKey) {
+      throw new Error("Model key must be specified for multi-service");
+    }
+    const item = this.services.get(modelKey);
+    if (!item) {
+      throw new Error(`No service found for model key: ${modelKey}`);
+    }
+    const service = item.service;
+    const model = item.useDefaultModel ? req.model : modelKey;
+    return await service.chat({ model, ...req }, options);
+  }
+  /**
+   * Delegates the embed call to the service matching the provided embed model key.
+   */
+  async embed(req, options) {
+    const modelKey = req.embedModel;
+    if (!modelKey) {
+      throw new Error("Embed model key must be specified for multi-service");
+    }
+    const item = this.services.get(modelKey);
+    if (!item) {
+      throw new Error(`No service found for embed model key: ${modelKey}`);
+    }
+    const service = item.service;
+    const embedModel = item.useDefaultModel ? req.embedModel : modelKey;
+    return await service.embed({ embedModel, ...req }, options);
+  }
+  /**
+   * Returns a composite ID built from the IDs of the underlying services.
+   */
+  getId() {
+    return "MultiServiceRouter:" + Array.from(this.services.values()).map((s) => s.service.getId()).join(",");
+  }
+  /**
+   * Returns the name of this router.
+   */
+  getName() {
+    return "MultiServiceRouter";
+  }
+  /**
+   * Aggregates all available models across the underlying services.
+   */
+  getModelList() {
+    return Array.from(this.services).filter(([, value]) => !value.isInternal).map(([key, { description, model }]) => ({
+      key,
+      description,
+      model
+    }));
+  }
+  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.
+   */
+  getFeatures(model) {
+    if (model) {
+      const service = this.services.get(model);
+      if (service) {
+        return service.service.getFeatures(model);
+      }
+    }
+    return { functions: false, streaming: false };
+  }
+  /**
+   * Returns aggregated metrics from the underlying service.
+   * Uses the metrics from the last service that was used,
+   * or falls back to the first service if none has been used.
+   */
+  getMetrics() {
+    const service = this.services.values().next().value;
+    if (!service) {
+      throw new Error("No service available to get metrics.");
+    }
+    return service.service.getMetrics();
+  }
+  /**
+   * Sets options on all underlying services.
+   */
+  setOptions(options) {
+    for (const service of this.services.values()) {
+      service.service.setOptions(options);
+    }
+  }
+  /**
+   * Returns the options from the last used service,
+   * or falls back to the first service if none has been used.
+   */
+  getOptions() {
+    const service = this.services.values().next().value;
+    if (!service) {
+      throw new Error("No service available to get options.");
+    }
+    return service.service.getOptions();
+  }
+};
+
 // prompts/rag.ts
 var AxRAG = class extends AxChainOfThought {
   genQuery;
@@ -8565,14 +8752,15 @@ export {
   AxLLMRequestTypeValues,
   AxMemory,
   AxMockAIService,
+  AxMultiServiceRouter,
   AxProgram,
   AxProgramWithSignature,
   AxPromptTemplate,
   AxRAG,
   AxRateLimiterTokenUsage,
-  AxRoute,
-  AxRouter,
   AxSignature,
+  AxSimpleClassifier,
+  AxSimpleClassifierClass,
   AxSpanKindValues,
   AxTestPrompt
 };