@ax-llm/ax 10.0.50 → 11.0.1

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,12 @@ 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;
+    this.id = crypto.randomUUID();
+    const model = this.models?.find((v) => v.key === defaults.model)?.model ?? defaults.model;
+    const embedModel = 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,13 +662,14 @@ var AxBaseAI = class {
   rt;
   fetch;
   tracer;
-  modelMap;
+  models;
   modelInfo;
   modelUsage;
   embedModelUsage;
-  models;
+  defaults;
   apiURL;
   name;
+  id;
   headers;
   supportFor;
   // Add private metrics tracking properties
@@ -703,6 +704,9 @@ var AxBaseAI = class {
   setName(name) {
     this.name = name;
   }
+  getId() {
+    return this.id;
+  }
   setAPIURL(apiURL) {
     this.apiURL = apiURL;
   }
@@ -733,9 +737,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 +747,26 @@ var AxBaseAI = class {
     };
   }
   getEmbedModelInfo() {
-    if (!this.models.embedModel) {
+    if (!this.defaults.embedModel) {
       return;
     }
     const mi = getModelInfo({
-      model: this.models.embedModel,
-      modelInfo: this.modelInfo,
-      modelMap: this.modelMap
+      model: this.defaults.embedModel,
+      modelInfo: this.modelInfo
     });
     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 +814,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 +958,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.defaults.embedModel;
     if (!embedModel) {
       throw new Error("No embed model defined");
     }
@@ -966,7 +969,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]: embedModel
           }
         },
         async (span) => {
@@ -1145,13 +1148,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
@@ -1439,7 +1442,7 @@ var AxAIAnthropic = class extends AxBaseAI {
     region,
     config,
     options,
-    modelMap
+    models
   }) {
     const isVertex = projectId !== void 0 && region !== void 0;
     let apiURL;
@@ -1475,10 +1478,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
     });
   }
 };
@@ -1955,51 +1958,70 @@ function createMessages2(req) {
     }
   });
 }
-var AxAIOpenAI = class extends AxBaseAI {
+var AxAIOpenAIBase = class extends AxBaseAI {
   constructor({
     apiKey,
     config,
     options,
     apiURL,
-    modelInfo = axModelInfoOpenAI,
-    modelMap
+    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",
       apiURL: apiURL ? apiURL : "https://api.openai.com/v1",
       headers: async () => ({ Authorization: `Bearer ${apiKey}` }),
       modelInfo,
-      models: {
-        model: _config.model,
-        embedModel: _config.embedModel
+      defaults: {
+        model: config.model,
+        embedModel: config.embedModel
       },
       options,
       supportFor: () => {
         return { functions: true, streaming: true };
       },
-      modelMap
+      models
     });
   }
 };
 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,
@@ -2007,7 +2029,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 +2044,13 @@ var AxAIAzureOpenAI = class extends AxAIOpenAI {
       ...axAIAzureOpenAIDefaultConfig(),
       ...config
     };
-    super({ apiKey, config: _config, options, modelMap });
+    super({
+      apiKey,
+      config: _config,
+      options,
+      models,
+      modelInfo: axModelInfoOpenAI
+    });
     const host = resourceName.includes("://") ? resourceName : `https://${resourceName}.openai.azure.com/`;
     super.setName("Azure OpenAI");
     super.setAPIURL(
@@ -2278,7 +2306,7 @@ var AxAICohere = class extends AxBaseAI {
     apiKey,
     config,
     options,
-    modelMap
+    models
   }) {
     if (!apiKey || apiKey === "") {
       throw new Error("Cohere API key not set");
@@ -2293,10 +2321,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
     });
   }
 };
@@ -2385,12 +2413,12 @@ var axAIDeepSeekDefaultConfig = () => structuredClone({
   model: "deepseek-chat" /* DeepSeekChat */,
   ...axBaseAIDefaultConfig()
 });
-var AxAIDeepSeek = class extends AxAIOpenAI {
+var AxAIDeepSeek = class extends AxAIOpenAIBase {
   constructor({
     apiKey,
     config,
     options,
-    modelMap
+    models
   }) {
     if (!apiKey || apiKey === "") {
       throw new Error("DeepSeek API key not set");
@@ -2405,7 +2433,7 @@ var AxAIDeepSeek = class extends AxAIOpenAI {
       options,
       apiURL: "https://api.deepseek.com",
       modelInfo: axModelInfoDeepSeek,
-      modelMap
+      models
     });
     super.setName("DeepSeek");
   }
@@ -2413,6 +2441,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 +2476,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 +2842,7 @@ var AxAIGoogleGemini = class extends AxBaseAI {
     region,
     config,
     options,
-    modelMap
+    models
   }) {
     const isVertex = projectId !== void 0 && region !== void 0;
     let apiURL;
@@ -2814,13 +2874,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
     });
   }
 };
@@ -2917,12 +2977,12 @@ 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,
     options,
-    modelMap
+    models
   }) {
     if (!apiKey || apiKey === "") {
       throw new Error("Groq API key not set");
@@ -2941,7 +3001,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 +3125,7 @@ var AxAIHuggingFace = class extends AxBaseAI {
     apiKey,
     config,
     options,
-    modelMap
+    models
   }) {
     if (!apiKey || apiKey === "") {
       throw new Error("HuggingFace API key not set");
@@ -3080,10 +3140,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
     });
   }
 };
@@ -3162,12 +3222,12 @@ var axAIMistralDefaultConfig = () => structuredClone({
   model: "mistral-small-latest" /* MistralSmall */,
   ...axBaseAIDefaultConfig()
 });
-var AxAIMistral = class extends AxAIOpenAI {
+var AxAIMistral = class extends AxAIOpenAIBase {
   constructor({
     apiKey,
     config,
     options,
-    modelMap
+    models
   }) {
     if (!apiKey || apiKey === "") {
       throw new Error("Mistral API key not set");
@@ -3182,7 +3242,7 @@ var AxAIMistral = class extends AxAIOpenAI {
       options,
       apiURL: "https://api.mistral.ai/v1",
       modelInfo: axModelInfoMistral,
-      modelMap
+      models
     });
     super.setName("Mistral");
   }
@@ -3194,13 +3254,13 @@ 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",
     config,
     options,
-    modelMap
+    models
   }) {
     const _config = {
       ...axAIOllamaDefaultConfig(),
@@ -3211,7 +3271,8 @@ var AxAIOllama = class extends AxAIOpenAI {
       options,
       config: _config,
       apiURL: url,
-      modelMap
+      models,
+      modelInfo: []
     });
     super.setName("Ollama");
   }
@@ -3414,7 +3475,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 +3490,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
     });
   }
 };
@@ -3448,12 +3509,12 @@ var axAITogetherDefaultConfig = () => structuredClone({
   model: "mistralai/Mixtral-8x7B-Instruct-v0.1",
   ...axBaseAIDefaultConfig()
 });
-var AxAITogether = class extends AxAIOpenAI {
+var AxAITogether = class extends AxAIOpenAIBase {
   constructor({
     apiKey,
     config,
     options,
-    modelMap
+    models
   }) {
     if (!apiKey || apiKey === "") {
       throw new Error("Together API key not set");
@@ -3468,7 +3529,7 @@ var AxAITogether = class extends AxAIOpenAI {
       options,
       apiURL: "https://api.together.xyz/v1",
       modelInfo: axModelInfoTogether,
-      modelMap
+      models
     });
     super.setName("Together");
   }
@@ -3522,6 +3583,9 @@ var AxAI = class {
   getName() {
     return this.ai.getName();
   }
+  getId() {
+    return this.ai.getId();
+  }
   getModelInfo() {
     return this.ai.getModelInfo();
   }
@@ -3531,8 +3595,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 +5729,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 +5990,7 @@ var AxAgent = class {
   ai;
   signature;
   program;
+  functions;
   agents;
   name;
   description;
@@ -5941,14 +6006,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 +6019,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 +6032,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 +6060,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 +6108,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";
@@ -6094,6 +6187,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.");
@@ -6107,6 +6204,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.
@@ -6122,8 +6222,8 @@ var AxBalancer = class _AxBalancer {
     const bTotalCost = (bInfo.promptTokenCostPer1M || Infinity) + (bInfo.completionTokenCostPer1M || Infinity);
     return aTotalCost - bTotalCost;
   };
-  getModelMap() {
-    throw new Error("Method not implemented.");
+  getModelList() {
+    return this.currentService.getModelList();
   }
   getNextService() {
     const cs = this.services[++this.currentServiceIndex];
@@ -6144,6 +6244,9 @@ var AxBalancer = class _AxBalancer {
   getName() {
     return this.currentService.getName();
   }
+  getId() {
+    return this.currentService.getId();
+  }
   getModelInfo() {
     return this.currentService.getModelInfo();
   }
@@ -6156,11 +6259,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;
@@ -6181,38 +6329,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()}`);
-        }
       }
     }
   }
@@ -7705,6 +7844,7 @@ var AxEmbeddingAdapter = class {
 var AxMockAIService = class {
   constructor(config = {}) {
     this.config = config;
+    this.config.id = this.config.id ?? crypto.randomUUID();
   }
   options = {};
   metrics = {
@@ -7720,6 +7860,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",
@@ -7739,8 +7882,8 @@ var AxMockAIService = class {
       streaming: this.config.features?.streaming ?? false
     };
   }
-  getModelMap() {
-    return this.config.modelMap;
+  getModelList() {
+    return this.config.models;
   }
   getMetrics() {
     return this.metrics;
@@ -7875,6 +8018,7 @@ export {
   AxAIMistralModel,
   AxAIOllama,
   AxAIOpenAI,
+  AxAIOpenAIBase,
   AxAIOpenAIEmbedModel,
   AxAIOpenAIModel,
   AxAIReka,