@ax-llm/ax 10.0.50 → 11.0.0

package/index.js

@@ -5,9 +5,9 @@ import { SpanKind } from "@opentelemetry/api";
 function getModelInfo({
   model,
   modelInfo,
-  modelMap = {}
+  models
 }) {
-  const mappedModel = modelMap?.[model] ?? model;
+  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+$/, "");
@@ -636,10 +636,10 @@ var AxBaseAI = class {
     apiURL,
     headers,
     modelInfo,
-    models,
+    defaults,
     options = {},
     supportFor,
-    modelMap
+    models
   }) {
     this.aiImpl = aiImpl;
     this.name = name;
@@ -648,12 +648,11 @@ var AxBaseAI = class {
     this.supportFor = supportFor;
     this.tracer = options.tracer;
     this.modelInfo = modelInfo;
-    this.modelMap = modelMap;
-    this.models = {
-      model: modelMap?.[models.model] ?? models.model,
-      embedModel: modelMap?.[models.embedModel ?? ""] ?? models.embedModel
-    };
-    if (!models.model || typeof models.model !== "string" || models.model === "") {
+    this.models = models;
+    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;
+    this.defaults = { model, embedModel };
+    if (!defaults.model || typeof defaults.model !== "string" || defaults.model === "") {
       throw new Error("No model defined");
     }
     this.setOptions(options);
@@ -662,11 +661,11 @@ var AxBaseAI = class {
   rt;
   fetch;
   tracer;
-  modelMap;
+  models;
   modelInfo;
   modelUsage;
   embedModelUsage;
-  models;
+  defaults;
   apiURL;
   name;
   headers;
@@ -733,9 +732,9 @@ var AxBaseAI = class {
   }
   getModelInfo() {
     const mi = getModelInfo({
-      model: this.models.model,
+      model: this.defaults.model,
       modelInfo: this.modelInfo,
-      modelMap: this.modelMap
+      models: this.models
     });
     return {
       ...mi,
@@ -743,27 +742,27 @@ var AxBaseAI = class {
     };
   }
   getEmbedModelInfo() {
-    if (!this.models.embedModel) {
+    if (!this.defaults.embedModel) {
       return;
     }
     const mi = getModelInfo({
-      model: this.models.embedModel,
+      model: this.defaults.embedModel,
       modelInfo: this.modelInfo,
-      modelMap: this.modelMap
+      models: this.models
     });
     return {
       ...mi,
       provider: this.name
     };
   }
-  getModelMap() {
-    return this.modelMap;
+  getModelList() {
+    return this.models;
   }
   getName() {
     return this.name;
   }
   getFeatures(model) {
-    return typeof this.supportFor === "function" ? this.supportFor(model ?? this.models.model) : this.supportFor;
+    return typeof this.supportFor === "function" ? this.supportFor(model ?? this.defaults.model) : this.supportFor;
   }
   // Method to calculate percentiles
   calculatePercentile(samples, percentile) {
@@ -811,7 +810,7 @@ var AxBaseAI = class {
     }
   }
   async _chat1(req, options) {
-    const model = req.model ? this.modelMap?.[req.model] ?? req.model : this.modelMap?.[this.models.model] ?? this.models.model;
+    const model = req.model ? this.models?.find((v) => v.key === req.model)?.model ?? req.model : this.defaults.model;
     const modelConfig = {
       ...this.aiImpl.getModelConfig(),
       ...req.modelConfig
@@ -955,7 +954,7 @@ var AxBaseAI = class {
     }
   }
   async _embed1(req, options) {
-    const embedModel = req.embedModel ? this.modelMap?.[req.embedModel] ?? req.embedModel : this.modelMap?.[this.models.embedModel ?? ""] ?? this.models.embedModel;
+    const embedModel = req.embedModel ? this.models?.find((v) => v.key === req.embedModel)?.model ?? req.embedModel : this.defaults.embedModel;
     if (!embedModel) {
       throw new Error("No embed model defined");
     }
@@ -966,7 +965,7 @@ var AxBaseAI = class {
           kind: SpanKind.SERVER,
           attributes: {
             [axSpanAttributes.LLM_SYSTEM]: this.name,
-            [axSpanAttributes.LLM_REQUEST_MODEL]: req.embedModel ?? this.models.embedModel
+            [axSpanAttributes.LLM_REQUEST_MODEL]: req.embedModel ?? this.defaults.embedModel
           }
         },
         async (span) => {
@@ -1439,7 +1438,7 @@ var AxAIAnthropic = class extends AxBaseAI {
     region,
     config,
     options,
-    modelMap
+    models
   }) {
     const isVertex = projectId !== void 0 && region !== void 0;
     let apiURL;
@@ -1475,10 +1474,10 @@ var AxAIAnthropic = class extends AxBaseAI {
       apiURL,
       headers,
       modelInfo: axModelInfoAnthropic,
-      models: { model: _config.model },
+      defaults: { model: _config.model },
       options,
-      supportFor: { functions: true, streaming: true },
-      modelMap
+      supportFor: { functions: true, streaming: true, functionCot: true },
+      models
     });
   }
 };
@@ -1962,7 +1961,7 @@ var AxAIOpenAI = class extends AxBaseAI {
     options,
     apiURL,
     modelInfo = axModelInfoOpenAI,
-    modelMap
+    models
   }) {
     if (!apiKey || apiKey === "") {
       throw new Error("OpenAI API key not set");
@@ -1981,7 +1980,7 @@ var AxAIOpenAI = class extends AxBaseAI {
       apiURL: apiURL ? apiURL : "https://api.openai.com/v1",
       headers: async () => ({ Authorization: `Bearer ${apiKey}` }),
       modelInfo,
-      models: {
+      defaults: {
         model: _config.model,
         embedModel: _config.embedModel
       },
@@ -1989,7 +1988,7 @@ var AxAIOpenAI = class extends AxBaseAI {
       supportFor: () => {
         return { functions: true, streaming: true };
       },
-      modelMap
+      models
     });
   }
 };
@@ -2007,7 +2006,7 @@ var AxAIAzureOpenAI = class extends AxAIOpenAI {
     version = "api-version=2024-02-15-preview",
     config,
     options,
-    modelMap
+    models
   }) {
     if (!apiKey || apiKey === "") {
       throw new Error("Azure OpenAPI API key not set");
@@ -2022,7 +2021,7 @@ var AxAIAzureOpenAI = class extends AxAIOpenAI {
       ...axAIAzureOpenAIDefaultConfig(),
       ...config
     };
-    super({ apiKey, config: _config, options, modelMap });
+    super({ apiKey, config: _config, options, models });
     const host = resourceName.includes("://") ? resourceName : `https://${resourceName}.openai.azure.com/`;
     super.setName("Azure OpenAI");
     super.setAPIURL(
@@ -2278,7 +2277,7 @@ var AxAICohere = class extends AxBaseAI {
     apiKey,
     config,
     options,
-    modelMap
+    models
   }) {
     if (!apiKey || apiKey === "") {
       throw new Error("Cohere API key not set");
@@ -2293,10 +2292,10 @@ var AxAICohere = class extends AxBaseAI {
       apiURL: "https://api.cohere.ai/v1",
       headers: async () => ({ Authorization: `Bearer ${apiKey}` }),
       modelInfo: axModelInfoCohere,
-      models: { model: _config.model },
+      defaults: { model: _config.model },
       supportFor: { functions: true, streaming: true },
       options,
-      modelMap
+      models
     });
   }
 };
@@ -2390,7 +2389,7 @@ var AxAIDeepSeek = class extends AxAIOpenAI {
     apiKey,
     config,
     options,
-    modelMap
+    models
   }) {
     if (!apiKey || apiKey === "") {
       throw new Error("DeepSeek API key not set");
@@ -2405,7 +2404,7 @@ var AxAIDeepSeek = class extends AxAIOpenAI {
       options,
       apiURL: "https://api.deepseek.com",
       modelInfo: axModelInfoDeepSeek,
-      modelMap
+      models
     });
     super.setName("DeepSeek");
   }
@@ -2413,6 +2412,10 @@ var AxAIDeepSeek = class extends AxAIOpenAI {
 
 // ai/google-gemini/types.ts
 var AxAIGoogleGeminiModel = /* @__PURE__ */ ((AxAIGoogleGeminiModel2) => {
+  AxAIGoogleGeminiModel2["Gemini20Pro"] = "gemini-2.0-pro-exp-02-05";
+  AxAIGoogleGeminiModel2["Gemini20Flash"] = "gemini-2.0-flash";
+  AxAIGoogleGeminiModel2["Gemini20FlashLite"] = "gemini-2.0-flash-lite-preview-02-05";
+  AxAIGoogleGeminiModel2["Gemini20FlashThinking"] = "gemini-2.0-flash-thinking-exp-01-21";
   AxAIGoogleGeminiModel2["Gemini1Pro"] = "gemini-1.0-pro";
   AxAIGoogleGeminiModel2["Gemini15Flash"] = "gemini-1.5-flash";
   AxAIGoogleGeminiModel2["Gemini15Flash8B"] = "gemini-1.5-flash-8b";
@@ -2444,31 +2447,59 @@ var AxAIGoogleGeminiSafetyThreshold = /* @__PURE__ */ ((AxAIGoogleGeminiSafetyTh
 
 // ai/google-gemini/info.ts
 var axModelInfoGoogleGemini = [
+  {
+    name: "gemini-2.0-pro-exp-02-05" /* Gemini20Pro */,
+    currency: "usd",
+    characterIsToken: false,
+    promptTokenCostPer1M: 0,
+    completionTokenCostPer1M: 0
+  },
+  {
+    name: "gemini-2.0-flash" /* Gemini20Flash */,
+    currency: "usd",
+    characterIsToken: false,
+    promptTokenCostPer1M: 0.01,
+    completionTokenCostPer1M: 0.4
+  },
+  {
+    name: "gemini-2.0-flash-lite-preview-02-05" /* Gemini20FlashLite */,
+    currency: "usd",
+    characterIsToken: false,
+    promptTokenCostPer1M: 0,
+    completionTokenCostPer1M: 0
+  },
+  {
+    name: "gemini-2.0-flash-thinking-exp-01-21" /* Gemini20FlashThinking */,
+    currency: "usd",
+    characterIsToken: false,
+    promptTokenCostPer1M: 0,
+    completionTokenCostPer1M: 0
+  },
   {
     name: "gemini-1.5-flash" /* Gemini15Flash */,
     currency: "usd",
-    characterIsToken: true,
+    characterIsToken: false,
     promptTokenCostPer1M: 0.075,
     completionTokenCostPer1M: 0.3
   },
   {
     name: "gemini-1.5-flash-8b" /* Gemini15Flash8B */,
     currency: "usd",
-    characterIsToken: true,
+    characterIsToken: false,
     promptTokenCostPer1M: 0.0375,
     completionTokenCostPer1M: 0.15
   },
   {
     name: "gemini-1.5-pro" /* Gemini15Pro */,
     currency: "usd",
-    characterIsToken: true,
+    characterIsToken: false,
     promptTokenCostPer1M: 1.25,
     completionTokenCostPer1M: 5
   },
   {
     name: "gemini-1.0-pro" /* Gemini1Pro */,
     currency: "usd",
-    characterIsToken: true,
+    characterIsToken: false,
     promptTokenCostPer1M: 0.5,
     completionTokenCostPer1M: 1.5
   }
@@ -2782,7 +2813,7 @@ var AxAIGoogleGemini = class extends AxBaseAI {
     region,
     config,
     options,
-    modelMap
+    models
   }) {
     const isVertex = projectId !== void 0 && region !== void 0;
     let apiURL;
@@ -2814,13 +2845,13 @@ var AxAIGoogleGemini = class extends AxBaseAI {
       apiURL,
       headers,
       modelInfo: axModelInfoGoogleGemini,
-      models: {
+      defaults: {
         model: _config.model,
         embedModel: _config.embedModel
       },
       options,
       supportFor: { functions: true, streaming: true },
-      modelMap
+      models
     });
   }
 };
@@ -2922,7 +2953,7 @@ var AxAIGroq = class extends AxAIOpenAI {
     apiKey,
     config,
     options,
-    modelMap
+    models
   }) {
     if (!apiKey || apiKey === "") {
       throw new Error("Groq API key not set");
@@ -2941,7 +2972,7 @@ var AxAIGroq = class extends AxAIOpenAI {
       options: _options,
       modelInfo: axModelInfoGroq,
       apiURL: "https://api.groq.com/openai/v1",
-      modelMap
+      models
     });
     super.setName("Groq");
     this.setOptions(_options);
@@ -3065,7 +3096,7 @@ var AxAIHuggingFace = class extends AxBaseAI {
     apiKey,
     config,
     options,
-    modelMap
+    models
   }) {
     if (!apiKey || apiKey === "") {
       throw new Error("HuggingFace API key not set");
@@ -3080,10 +3111,10 @@ var AxAIHuggingFace = class extends AxBaseAI {
       apiURL: "https://api-inference.huggingface.co",
       headers: async () => ({ Authorization: `Bearer ${apiKey}` }),
       modelInfo: axModelInfoHuggingFace,
-      models: { model: _config.model },
+      defaults: { model: _config.model },
       options,
       supportFor: { functions: false, streaming: false },
-      modelMap
+      models
     });
   }
 };
@@ -3167,7 +3198,7 @@ var AxAIMistral = class extends AxAIOpenAI {
     apiKey,
     config,
     options,
-    modelMap
+    models
   }) {
     if (!apiKey || apiKey === "") {
       throw new Error("Mistral API key not set");
@@ -3182,7 +3213,7 @@ var AxAIMistral = class extends AxAIOpenAI {
       options,
       apiURL: "https://api.mistral.ai/v1",
       modelInfo: axModelInfoMistral,
-      modelMap
+      models
     });
     super.setName("Mistral");
   }
@@ -3200,7 +3231,7 @@ var AxAIOllama = class extends AxAIOpenAI {
     url = "http://localhost:11434/v1",
     config,
     options,
-    modelMap
+    models
   }) {
     const _config = {
       ...axAIOllamaDefaultConfig(),
@@ -3211,7 +3242,7 @@ var AxAIOllama = class extends AxAIOpenAI {
       options,
       config: _config,
       apiURL: url,
-      modelMap
+      models
     });
     super.setName("Ollama");
   }
@@ -3414,7 +3445,7 @@ var AxAIReka = class extends AxBaseAI {
     options,
     apiURL,
     modelInfo = axModelInfoReka,
-    modelMap
+    models
   }) {
     if (!apiKey || apiKey === "") {
       throw new Error("Reka API key not set");
@@ -3429,12 +3460,12 @@ var AxAIReka = class extends AxBaseAI {
       apiURL: apiURL ? apiURL : "https://api.reka.ai/v1/chat",
       headers: async () => ({ "X-Api-Key": apiKey }),
       modelInfo,
-      models: {
+      defaults: {
         model: _config.model
       },
       options,
       supportFor: { functions: true, streaming: true },
-      modelMap
+      models
     });
   }
 };
@@ -3453,7 +3484,7 @@ var AxAITogether = class extends AxAIOpenAI {
     apiKey,
     config,
     options,
-    modelMap
+    models
   }) {
     if (!apiKey || apiKey === "") {
       throw new Error("Together API key not set");
@@ -3468,7 +3499,7 @@ var AxAITogether = class extends AxAIOpenAI {
       options,
       apiURL: "https://api.together.xyz/v1",
       modelInfo: axModelInfoTogether,
-      modelMap
+      models
     });
     super.setName("Together");
   }
@@ -3531,8 +3562,8 @@ var AxAI = class {
   getFeatures(model) {
     return this.ai.getFeatures(model);
   }
-  getModelMap() {
-    return this.ai.getModelMap();
+  getModelList() {
+    return this.ai.getModelList();
   }
   getMetrics() {
     return this.ai.getMetrics();
@@ -5665,7 +5696,7 @@ var AxGen = class extends AxProgramWithSignature {
     traceId,
     functions
   }) {
-    const streamingValidation = !functions || functions.length == 0;
+    const streamingValidation = ai.getFeatures().functionCot !== true;
     const functionCalls = [];
     const values = {};
     const xstate = {
@@ -5926,6 +5957,7 @@ var AxAgent = class {
   ai;
   signature;
   program;
+  functions;
   agents;
   name;
   description;
@@ -5941,14 +5973,9 @@ var AxAgent = class {
   }, options) {
     this.ai = ai;
     this.agents = agents;
+    this.functions = functions;
     this.signature = new AxSignature(signature);
     this.signature.setDescription(description);
-    const funcs = [
-      ...functions ?? [],
-      ...agents?.map((a) => a.getFunction()) ?? []
-    ];
-    const opt = { ...options, functions: funcs };
-    this.program = new AxGen(this.signature, opt);
     if (!name || name.length < 5) {
       throw new Error(
         `Agent name must be at least 10 characters (more descriptive): ${name}`
@@ -5959,6 +5986,10 @@ var AxAgent = class {
         `Agent description must be at least 20 characters (explain in detail what the agent does): ${description}`
       );
     }
+    this.program = new AxGen(this.signature, options);
+    for (const agent of agents ?? []) {
+      this.program.register(agent);
+    }
     this.name = name;
     this.description = description;
     this.subAgentList = agents?.map((a) => a.getFunction().name).join(", ");
@@ -5968,8 +5999,9 @@ var AxAgent = class {
       parameters: this.signature.toJSONSchema(),
       func: () => this.forward
     };
-    for (const agent of agents ?? []) {
-      this.program.register(agent);
+    const mm = ai?.getModelList();
+    if (mm) {
+      this.func.parameters = addModelParameter(this.func.parameters, mm);
     }
   }
   setExamples(examples) {
@@ -5995,38 +6027,41 @@ var AxAgent = class {
   }
   getFunction() {
     const boundFunc = this.forward.bind(this);
-    const wrappedFunc = (values, options) => {
+    const wrappedFunc = (valuesAndModel, options) => {
+      const { model, ...values } = valuesAndModel;
       const ai = this.ai ?? options?.ai;
       if (!ai) {
         throw new Error("AI service is required to run the agent");
       }
-      return boundFunc(ai, values, options);
+      return boundFunc(ai, values, { ...options, model });
     };
     return {
       ...this.func,
       func: wrappedFunc
     };
   }
-  init(ai, options) {
-    const _ai = this.ai ?? ai;
-    const funcs = [
-      ...options?.functions ?? [],
-      ...this.agents?.map((a) => a.getFunction()) ?? []
+  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
+    );
+    const functions = [
+      ...options?.functions ?? this.functions ?? [],
+      ...agentFuncs ?? []
     ];
-    const opt = options;
-    if (funcs.length > 0) {
-      const opt2 = { ...options, functions: funcs };
-      this.program = new AxGen(this.signature, opt2);
-    }
-    return { _ai, opt };
+    return { ai, functions };
   }
-  async forward(ai, values, options) {
-    const { _ai, opt } = this.init(ai, options);
-    return await this.program.forward(_ai, values, opt);
+  async forward(parentAi, values, options) {
+    const { ai, functions } = this.init(parentAi, options);
+    return await this.program.forward(ai, values, { ...options, functions });
   }
-  async *streamingForward(ai, values, options) {
-    const { _ai, opt } = this.init(ai, options);
-    return yield* this.program.streamingForward(_ai, values, opt);
+  async *streamingForward(parentAi, values, options) {
+    const { ai, functions } = this.init(parentAi, options);
+    return yield* this.program.streamingForward(ai, values, {
+      ...options,
+      functions
+    });
   }
 };
 function toCamelCase(inputString) {
@@ -6040,6 +6075,31 @@ function toCamelCase(inputString) {
   }).join("");
   return camelCaseString;
 }
+function addModelParameter(parameters, models) {
+  const baseSchema = parameters ? structuredClone(parameters) : {
+    type: "object",
+    properties: {},
+    required: []
+  };
+  if (baseSchema.properties?.model) {
+    return baseSchema;
+  }
+  const modelProperty = {
+    type: "string",
+    enum: models.map((m) => m.key),
+    description: `The AI model to use for this function call. Available options: ${models.map((m) => `${m.key}: ${m.description}`).join(" | ")}`
+  };
+  const newProperties = {
+    ...baseSchema.properties ?? {},
+    model: modelProperty
+  };
+  const newRequired = [...baseSchema.required ?? [], "model"];
+  return {
+    ...baseSchema,
+    properties: newProperties,
+    required: newRequired
+  };
+}
 
 // docs/tika.ts
 import { createReadStream } from "node:fs";
@@ -6122,7 +6182,7 @@ var AxBalancer = class _AxBalancer {
     const bTotalCost = (bInfo.promptTokenCostPer1M || Infinity) + (bInfo.completionTokenCostPer1M || Infinity);
     return aTotalCost - bTotalCost;
   };
-  getModelMap() {
+  getModelList() {
     throw new Error("Method not implemented.");
   }
   getNextService() {
@@ -7739,8 +7799,8 @@ var AxMockAIService = class {
       streaming: this.config.features?.streaming ?? false
     };
   }
-  getModelMap() {
-    return this.config.modelMap;
+  getModelList() {
+    return this.config.models;
   }
   getMetrics() {
     return this.metrics;