@ax-llm/ax 11.0.47 → 11.0.49

package/index.cjs

@@ -365,14 +365,14 @@ var defaultRetryConfig = {
 var defaultTimeoutMs = 3e4;
 var textDecoderStream = import_web3.TextDecoderStream ?? TextDecoderStreamPolyfill;
 var AxAIServiceError = class extends Error {
-  constructor(message, url, requestBody, context = {}) {
+  constructor(message, url, requestBody, context3 = {}) {
     super(message);
     this.url = url;
     this.requestBody = requestBody;
     this.name = this.constructor.name;
     this.timestamp = (/* @__PURE__ */ new Date()).toISOString();
     this.errorId = crypto.randomUUID();
-    this.context = context;
+    this.context = context3;
     this.stack = this.toString();
   }
   timestamp;
@@ -394,11 +394,11 @@ var AxAIServiceError = class extends Error {
   }
 };
 var AxAIServiceStatusError = class extends AxAIServiceError {
-  constructor(status, statusText, url, requestBody, context) {
+  constructor(status, statusText, url, requestBody, context3) {
     super(`HTTP ${status} - ${statusText}`, url, requestBody, {
       httpStatus: status,
       httpStatusText: statusText,
-      ...context
+      ...context3
     });
     this.status = status;
     this.statusText = statusText;
@@ -406,11 +406,11 @@ var AxAIServiceStatusError = class extends AxAIServiceError {
   }
 };
 var AxAIServiceNetworkError = class extends AxAIServiceError {
-  constructor(originalError, url, requestBody, context) {
+  constructor(originalError, url, requestBody, context3) {
     super(`Network Error: ${originalError.message}`, url, requestBody, {
       originalErrorName: originalError.name,
       originalErrorStack: originalError.stack,
-      ...context
+      ...context3
     });
     this.originalError = originalError;
     this.name = this.constructor.name;
@@ -418,33 +418,33 @@ var AxAIServiceNetworkError = class extends AxAIServiceError {
   }
 };
 var AxAIServiceResponseError = class extends AxAIServiceError {
-  constructor(message, url, requestBody, context) {
-    super(message, url, requestBody, context);
+  constructor(message, url, requestBody, context3) {
+    super(message, url, requestBody, context3);
     this.name = this.constructor.name;
   }
 };
 var AxAIServiceStreamTerminatedError = class extends AxAIServiceError {
-  constructor(url, requestBody, lastChunk, context) {
+  constructor(url, requestBody, lastChunk, context3) {
     super("Stream terminated unexpectedly by remote host", url, requestBody, {
       lastChunk,
-      ...context
+      ...context3
     });
     this.lastChunk = lastChunk;
     this.name = this.constructor.name;
   }
 };
 var AxAIServiceTimeoutError = class extends AxAIServiceError {
-  constructor(url, timeoutMs, requestBody, context) {
+  constructor(url, timeoutMs, requestBody, context3) {
     super(`Request timeout after ${timeoutMs}ms`, url, requestBody, {
       timeoutMs,
-      ...context
+      ...context3
     });
     this.name = this.constructor.name;
   }
 };
 var AxAIServiceAuthenticationError = class extends AxAIServiceError {
-  constructor(url, requestBody, context) {
-    super("Authentication failed", url, requestBody, context);
+  constructor(url, requestBody, context3) {
+    super("Authentication failed", url, requestBody, context3);
     this.name = this.constructor.name;
   }
 };
@@ -490,14 +490,12 @@ var apiCall = async (api, json) => {
       );
     }
   }
-  if (api.span?.isRecording()) {
-    api.span.setAttributes({
-      "http.request.method": api.put ? "PUT" : "POST",
-      "url.full": apiUrl.href,
-      "request.id": requestId,
-      "request.startTime": metrics.startTime
-    });
-  }
+  api.span?.setAttributes({
+    "http.request.method": api.put ? "PUT" : "POST",
+    "url.full": apiUrl.href,
+    "request.id": requestId,
+    "request.startTime": metrics.startTime
+  });
   let attempt = 0;
   while (true) {
     const controller = new AbortController();
@@ -524,16 +522,14 @@ var apiCall = async (api, json) => {
         const delay = calculateRetryDelay(attempt, retryConfig);
         attempt++;
         updateRetryMetrics(metrics);
-        if (api.span?.isRecording()) {
-          api.span.addEvent("retry", {
-            attempt,
-            delay,
-            status: res.status,
-            "metrics.startTime": metrics.startTime,
-            "metrics.retryCount": metrics.retryCount,
-            "metrics.lastRetryTime": metrics.lastRetryTime
-          });
-        }
+        api.span?.addEvent("retry", {
+          attempt,
+          delay,
+          status: res.status,
+          "metrics.startTime": metrics.startTime,
+          "metrics.retryCount": metrics.retryCount,
+          "metrics.lastRetryTime": metrics.lastRetryTime
+        });
         await new Promise((resolve) => setTimeout(resolve, delay));
         continue;
       }
@@ -559,12 +555,10 @@ var apiCall = async (api, json) => {
             );
           }
         }
-        if (api.span?.isRecording()) {
-          api.span.setAttributes({
-            "response.time": Date.now() - metrics.startTime,
-            "response.retries": metrics.retryCount
-          });
-        }
+        api.span?.setAttributes({
+          "response.time": Date.now() - metrics.startTime,
+          "response.retries": metrics.retryCount
+        });
         return resJson;
       }
       if (!res.body) {
@@ -584,15 +578,11 @@ var apiCall = async (api, json) => {
           metrics.streamChunks = chunkCount;
           metrics.lastChunkTime = Date.now();
           controller2.enqueue(chunk);
-        },
-        flush() {
-          if (api.span?.isRecording()) {
-            api.span.setAttributes({
-              "stream.chunks": chunkCount,
-              "stream.duration": Date.now() - metrics.startTime,
-              "response.retries": metrics.retryCount
-            });
-          }
+          api.span?.addEvent("stream.chunk", {
+            "stream.chunks": chunkCount,
+            "stream.duration": Date.now() - metrics.startTime,
+            "response.retries": metrics.retryCount
+          });
         }
       });
       let closed = false;
@@ -677,16 +667,14 @@ var apiCall = async (api, json) => {
         const delay = calculateRetryDelay(attempt, retryConfig);
         attempt++;
         updateRetryMetrics(metrics);
-        if (api.span?.isRecording()) {
-          api.span.addEvent("retry", {
-            attempt,
-            delay,
-            error: error.message,
-            "metrics.startTime": metrics.startTime,
-            "metrics.retryCount": metrics.retryCount,
-            "metrics.lastRetryTime": metrics.lastRetryTime
-          });
-        }
+        api.span?.addEvent("retry", {
+          attempt,
+          delay,
+          error: error.message,
+          "metrics.startTime": metrics.startTime,
+          "metrics.retryCount": metrics.retryCount,
+          "metrics.lastRetryTime": metrics.lastRetryTime
+        });
         await new Promise((resolve) => setTimeout(resolve, delay));
         continue;
       }
@@ -1085,7 +1073,7 @@ var AxBaseAI = class {
       modelConfig.stream = false;
     }
     if (this.tracer) {
-      return await this.tracer?.startActiveSpan(
+      return await this.tracer.startActiveSpan(
         "AI Chat Request",
         {
           kind: import_api2.SpanKind.SERVER,
@@ -1103,6 +1091,7 @@ var AxBaseAI = class {
             [axSpanAttributes.LLM_REQUEST_LLM_IS_STREAMING]: modelConfig.stream
           }
         },
+        options?.traceContext ?? import_api2.context.active(),
         async (span) => {
           return await this._chat2(model, modelConfig, req, options, span);
         }
@@ -1284,6 +1273,7 @@ var AxBaseAI = class {
             [axSpanAttributes.LLM_REQUEST_MODEL]: embedModel
           }
         },
+        options?.traceContext ?? import_api2.context.active(),
         async (span) => {
           try {
             return await this._embed2(embedModel, req, options, span);
@@ -1985,6 +1975,22 @@ function mapFinishReason(stopReason) {
   }
 }
 
+// dsp/modelinfo.ts
+function getModelInfo({
+  model,
+  modelInfo,
+  models
+}) {
+  const modelEntry = models?.find((v) => v.key === model);
+  const mappedModel = modelEntry && "model" in modelEntry ? modelEntry.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 null;
+}
+
 // ai/openai/types.ts
 var AxAIOpenAIModel = /* @__PURE__ */ ((AxAIOpenAIModel2) => {
   AxAIOpenAIModel2["O1"] = "o1";
@@ -2234,6 +2240,10 @@ var AxAIOpenAIImpl = class {
           break;
         case "high":
           reqValue.reasoning_effort = "high";
+          break;
+        case "highest":
+          reqValue.reasoning_effort = "high";
+          break;
       }
     }
     if (this.chatReqUpdater) {
@@ -2450,15 +2460,7 @@ var AxAIOpenAIBase = class extends AxBaseAI {
         embedModel: config.embedModel
       },
       options,
-      supportFor: supportFor ?? ((model) => {
-        const modelInf = modelInfo.find((m) => m.name === model);
-        return {
-          functions: true,
-          streaming: true,
-          hasThinkingBudget: modelInf?.hasThinkingBudget ?? false,
-          hasShowThoughts: modelInf?.hasShowThoughts ?? false
-        };
-      }),
+      supportFor,
       models
     });
   }
@@ -2468,18 +2470,24 @@ var AxAIOpenAI = class extends AxAIOpenAIBase {
     apiKey,
     config,
     options,
-    models
+    models,
+    modelInfo
   }) {
     if (!apiKey || apiKey === "") {
       throw new Error("OpenAI API key not set");
     }
-    const supportForFn = (model) => {
-      const modelInf = axModelInfoOpenAI.find((m) => m.name === model);
+    modelInfo = [...axModelInfoOpenAI, ...modelInfo ?? []];
+    const supportFor = (model) => {
+      const mi = getModelInfo({
+        model,
+        modelInfo,
+        models
+      });
       return {
         functions: true,
         streaming: true,
-        hasThinkingBudget: modelInf?.hasThinkingBudget ?? false,
-        hasShowThoughts: modelInf?.hasShowThoughts ?? false
+        hasThinkingBudget: mi?.hasThinkingBudget ?? false,
+        hasShowThoughts: mi?.hasShowThoughts ?? false
       };
     };
     super({
@@ -2489,9 +2497,9 @@ var AxAIOpenAI = class extends AxAIOpenAIBase {
         ...config
       },
       options,
-      modelInfo: axModelInfoOpenAI,
+      modelInfo,
       models,
-      supportFor: supportForFn
+      supportFor
     });
     super.setName("OpenAI");
   }
@@ -2510,7 +2518,8 @@ var AxAIAzureOpenAI = class extends AxAIOpenAIBase {
     version = "api-version=2024-02-15-preview",
     config,
     options,
-    models
+    models,
+    modelInfo
   }) {
     if (!apiKey || apiKey === "") {
       throw new Error("Azure OpenAPI API key not set");
@@ -2525,21 +2534,27 @@ var AxAIAzureOpenAI = class extends AxAIOpenAIBase {
       ...axAIAzureOpenAIDefaultConfig(),
       ...config
     };
+    modelInfo = [...axModelInfoOpenAI, ...modelInfo ?? []];
+    const supportFor = (model) => {
+      const mi = getModelInfo({
+        model,
+        modelInfo,
+        models
+      });
+      return {
+        functions: true,
+        streaming: true,
+        hasThinkingBudget: mi?.hasThinkingBudget ?? false,
+        hasShowThoughts: mi?.hasShowThoughts ?? false
+      };
+    };
     super({
       apiKey,
       config: _config,
       options,
       models,
-      modelInfo: axModelInfoOpenAI,
-      supportFor: (model) => {
-        const modelInf = axModelInfoOpenAI.find((m) => m.name === model);
-        return {
-          functions: true,
-          streaming: true,
-          hasThinkingBudget: modelInf?.hasThinkingBudget ?? false,
-          hasShowThoughts: modelInf?.hasShowThoughts ?? false
-        };
-      }
+      modelInfo,
+      supportFor
     });
     const host = resourceName.includes("://") ? resourceName : `https://${resourceName}.openai.azure.com/`;
     super.setName("Azure OpenAI");
@@ -2889,6 +2904,7 @@ function createToolCall(functionCalls) {
 var AxAIDeepSeekModel = /* @__PURE__ */ ((AxAIDeepSeekModel2) => {
   AxAIDeepSeekModel2["DeepSeekChat"] = "deepseek-chat";
   AxAIDeepSeekModel2["DeepSeekCoder"] = "deepseek-coder";
+  AxAIDeepSeekModel2["DeepSeekReasoner"] = "deepseek-reasoner";
   return AxAIDeepSeekModel2;
 })(AxAIDeepSeekModel || {});
 
@@ -2897,14 +2913,14 @@ var axModelInfoDeepSeek = [
   {
     name: "deepseek-chat" /* DeepSeekChat */,
     currency: "USD",
-    promptTokenCostPer1M: 0.14,
-    completionTokenCostPer1M: 0.28
+    promptTokenCostPer1M: 0.27,
+    completionTokenCostPer1M: 1.1
   },
   {
-    name: "deepseek-coder" /* DeepSeekCoder */,
+    name: "deepseek-reasoner" /* DeepSeekReasoner */,
     currency: "USD",
-    promptTokenCostPer1M: 0.14,
-    completionTokenCostPer1M: 0.28
+    promptTokenCostPer1M: 0.55,
+    completionTokenCostPer1M: 2.19
   }
 ];
 
@@ -2922,7 +2938,8 @@ var AxAIDeepSeek = class extends AxAIOpenAIBase {
     apiKey,
     config,
     options,
-    models
+    models,
+    modelInfo
   }) {
     if (!apiKey || apiKey === "") {
       throw new Error("DeepSeek API key not set");
@@ -2931,12 +2948,19 @@ var AxAIDeepSeek = class extends AxAIOpenAIBase {
       ...axAIDeepSeekDefaultConfig(),
       ...config
     };
+    modelInfo = [...axModelInfoDeepSeek, ...modelInfo ?? []];
     super({
       apiKey,
       config: _config,
       options,
       apiURL: "https://api.deepseek.com",
-      modelInfo: axModelInfoDeepSeek,
+      modelInfo,
+      supportFor: {
+        functions: true,
+        streaming: true,
+        hasThinkingBudget: false,
+        hasShowThoughts: false
+      },
       models
     });
     super.setName("DeepSeek");
@@ -2949,7 +2973,6 @@ var AxAIGoogleGeminiModel = /* @__PURE__ */ ((AxAIGoogleGeminiModel2) => {
   AxAIGoogleGeminiModel2["Gemini25Flash"] = "gemini-2.5-flash-preview-04-17";
   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["Gemini15Flash002"] = "gemini-1.5-flash-002";
@@ -3025,13 +3048,6 @@ var axModelInfoGoogleGemini = [
     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",
@@ -3234,6 +3250,9 @@ var AxAIGoogleGeminiImpl = class {
         }
       });
     }
+    if (this.options?.googleSearch) {
+      tools.push({ google_search: {} });
+    }
     if (this.options?.urlContext) {
       tools.push({ url_context: {} });
     }
@@ -3272,16 +3291,19 @@ var AxAIGoogleGeminiImpl = class {
     if (config.thinkingTokenBudget) {
       switch (config.thinkingTokenBudget) {
         case "minimal":
-          thinkingConfig.thinkingBudget = 0;
+          thinkingConfig.thinkingBudget = 200;
           break;
         case "low":
-          thinkingConfig.thinkingBudget = 1024;
+          thinkingConfig.thinkingBudget = 800;
           break;
         case "medium":
-          thinkingConfig.thinkingBudget = 4096;
+          thinkingConfig.thinkingBudget = 5e3;
           break;
         case "high":
-          thinkingConfig.thinkingBudget = 8192;
+          thinkingConfig.thinkingBudget = 1e4;
+          break;
+        case "highest":
+          thinkingConfig.thinkingBudget = 24500;
           break;
       }
     }
@@ -3435,7 +3457,8 @@ var AxAIGoogleGemini = class extends AxBaseAI {
     endpointId,
     config,
     options,
-    models
+    models,
+    modelInfo
   }) {
     const isVertex = projectId !== void 0 && region !== void 0;
     let apiURL;
@@ -3474,26 +3497,32 @@ var AxAIGoogleGemini = class extends AxBaseAI {
       apiKey,
       options
     );
+    modelInfo = [...axModelInfoGoogleGemini, ...modelInfo ?? []];
+    const supportFor = (model) => {
+      const mi = getModelInfo({
+        model,
+        modelInfo,
+        models
+      });
+      return {
+        functions: true,
+        streaming: true,
+        hasThinkingBudget: mi?.hasThinkingBudget ?? false,
+        hasShowThoughts: mi?.hasShowThoughts ?? false,
+        functionCot: false
+      };
+    };
     super(aiImpl, {
       name: "GoogleGeminiAI",
       apiURL,
       headers,
-      modelInfo: axModelInfoGoogleGemini,
+      modelInfo,
       defaults: {
         model: _config.model,
         embedModel: _config.embedModel
       },
       options,
-      supportFor: (model) => {
-        const modelInf = axModelInfoGoogleGemini.find((m) => m.name === model);
-        return {
-          functions: true,
-          streaming: true,
-          hasThinkingBudget: modelInf?.hasThinkingBudget ?? false,
-          hasShowThoughts: modelInf?.hasShowThoughts ?? false,
-          functionCot: false
-        };
-      },
+      supportFor,
       models
     });
   }
@@ -3596,7 +3625,8 @@ var AxAIGroq = class extends AxAIOpenAIBase {
     apiKey,
     config,
     options,
-    models
+    models,
+    modelInfo
   }) {
     if (!apiKey || apiKey === "") {
       throw new Error("Groq API key not set");
@@ -3609,13 +3639,21 @@ var AxAIGroq = class extends AxAIOpenAIBase {
       ...options,
       streamingUsage: false
     };
+    modelInfo = [...axModelInfoGroq, ...modelInfo ?? []];
+    const supportFor = {
+      functions: true,
+      streaming: true,
+      hasThinkingBudget: false,
+      hasShowThoughts: false
+    };
     super({
       apiKey,
       config: _config,
       options: _options,
-      modelInfo: axModelInfoGroq,
+      modelInfo,
       apiURL: "https://api.groq.com/openai/v1",
-      models
+      models,
+      supportFor
     });
     super.setName("Groq");
     this.setOptions(_options);
@@ -3853,7 +3891,8 @@ var AxAIMistral = class extends AxAIOpenAIBase {
     apiKey,
     config,
     options,
-    models
+    models,
+    modelInfo
   }) {
     if (!apiKey || apiKey === "") {
       throw new Error("Mistral API key not set");
@@ -3862,13 +3901,21 @@ var AxAIMistral = class extends AxAIOpenAIBase {
       ...axAIMistralDefaultConfig(),
       ...config
     };
+    modelInfo = [...axModelInfoMistral, ...modelInfo ?? []];
+    const supportFor = {
+      functions: true,
+      streaming: true,
+      hasThinkingBudget: false,
+      hasShowThoughts: false
+    };
     super({
       apiKey,
       config: _config,
       options,
       apiURL: "https://api.mistral.ai/v1",
-      modelInfo: axModelInfoMistral,
-      models
+      modelInfo,
+      models,
+      supportFor
     });
     super.setName("Mistral");
   }
@@ -3903,7 +3950,13 @@ var AxAIOllama = class extends AxAIOpenAIBase {
       config: _config,
       apiURL: url,
       models,
-      modelInfo: []
+      modelInfo: [],
+      supportFor: {
+        functions: true,
+        streaming: true,
+        hasThinkingBudget: false,
+        hasShowThoughts: false
+      }
     });
     super.setName("Ollama");
   }
@@ -4154,7 +4207,8 @@ var AxAITogether = class extends AxAIOpenAIBase {
     apiKey,
     config,
     options,
-    models
+    models,
+    modelInfo
   }) {
     if (!apiKey || apiKey === "") {
       throw new Error("Together API key not set");
@@ -4163,13 +4217,21 @@ var AxAITogether = class extends AxAIOpenAIBase {
       ...axAITogetherDefaultConfig(),
       ...config
     };
+    modelInfo = [...axModelInfoTogether, ...modelInfo ?? []];
+    const supportFor = {
+      functions: true,
+      streaming: true,
+      hasThinkingBudget: false,
+      hasShowThoughts: false
+    };
     super({
       apiKey,
       config: _config,
       options,
       apiURL: "https://api.together.xyz/v1",
-      modelInfo: axModelInfoTogether,
-      models
+      modelInfo,
+      models,
+      supportFor
     });
     super.setName("Together");
   }
@@ -4315,7 +4377,8 @@ var AxAIGrok = class extends AxAIOpenAIBase {
     apiKey,
     config,
     options,
-    models
+    models,
+    modelInfo
   }) {
     if (!apiKey || apiKey === "") {
       throw new Error("Grok API key not set");
@@ -4324,22 +4387,28 @@ var AxAIGrok = class extends AxAIOpenAIBase {
       ...axAIGrokDefaultConfig(),
       ...config
     };
+    modelInfo = [...axModelInfoGrok, ...modelInfo ?? []];
+    const supportFor = (model) => {
+      const mi = getModelInfo({
+        model,
+        modelInfo,
+        models
+      });
+      return {
+        functions: true,
+        streaming: true,
+        hasThinkingBudget: mi?.hasThinkingBudget ?? false,
+        hasShowThoughts: mi?.hasShowThoughts ?? false
+      };
+    };
     super({
       apiKey,
       config: _config,
       options,
       apiURL: "https://api.x.ai/v1",
-      modelInfo: axModelInfoGrok,
+      modelInfo,
       models,
-      supportFor: (model) => {
-        const modelInf = axModelInfoGrok.find((m) => m.name === model);
-        return {
-          functions: true,
-          streaming: true,
-          hasThinkingBudget: modelInf?.hasThinkingBudget ?? false,
-          hasShowThoughts: modelInf?.hasShowThoughts ?? false
-        };
-      }
+      supportFor
     });
     super.setName("Grok");
   }
@@ -4682,9 +4751,9 @@ var SignatureParser = class {
       };
     } catch (error) {
       const errorMessage = error instanceof Error ? error.message : "Unknown error";
-      const context = this.getErrorContext();
+      const context3 = this.getErrorContext();
       throw new Error(`${errorMessage}
-${context}`);
+${context3}`);
     }
   }
   getErrorContext() {
@@ -6783,7 +6852,9 @@ var AxGen = class extends AxProgramWithSignature {
   async forwardSendRequest({
     ai,
     mem,
-    options
+    options,
+    traceContext,
+    firstStep
   }) {
     const {
       sessionId,
@@ -6801,7 +6872,10 @@ var AxGen = class extends AxProgramWithSignature {
       throw new Error("No chat prompt found");
     }
     const functions = _functions?.map((f) => "toFunction" in f ? f.toFunction() : f)?.flat();
-    const functionCall = _functionCall ?? this.options?.functionCall;
+    let functionCall = _functionCall ?? this.options?.functionCall;
+    if (!firstStep && (functionCall === "required" || typeof functionCall === "function")) {
+      functionCall = void 0;
+    }
     const res = await ai.chat(
       {
         chatPrompt,
@@ -6816,7 +6890,8 @@ var AxGen = class extends AxProgramWithSignature {
         rateLimiter,
         stream,
         debug: false,
-        thinkingTokenBudget
+        thinkingTokenBudget,
+        traceContext
       }
     );
     return res;
@@ -6825,7 +6900,9 @@ var AxGen = class extends AxProgramWithSignature {
     ai,
     mem,
     options,
-    span
+    firstStep,
+    span,
+    traceContext
   }) {
     const { sessionId, traceId, functions: _functions } = options ?? {};
     const fastFail = options?.fastFail ?? this.options?.fastFail;
@@ -6834,7 +6911,9 @@ var AxGen = class extends AxProgramWithSignature {
     const res = await this.forwardSendRequest({
       ai,
       mem,
-      options
+      options,
+      traceContext,
+      firstStep
     });
     if (res instanceof import_web5.ReadableStream) {
       yield* this.processStreamingResponse({
@@ -7072,7 +7151,7 @@ Content: ${result.content}`
     }
     return { ...this.values };
   }
-  async *_forward2(ai, values, options, span) {
+  async *_forward2(ai, values, options, span, traceContext) {
     const stopFunction = (options?.stopFunction ?? this.options?.stopFunction)?.toLowerCase();
     const maxRetries = options.maxRetries ?? this.options?.maxRetries ?? 10;
     const maxSteps = options.maxSteps ?? this.options?.maxSteps ?? 10;
@@ -7094,9 +7173,17 @@ Content: ${result.content}`
     });
     mem.add(prompt, options?.sessionId);
     multiStepLoop: for (let n = 0; n < maxSteps; n++) {
+      const firstStep = n === 0;
       for (let errCount = 0; errCount < maxRetries; errCount++) {
         try {
-          const generator = this.forwardCore({ options, ai, mem, span });
+          const generator = this.forwardCore({
+            options,
+            ai,
+            mem,
+            firstStep,
+            span,
+            traceContext
+          });
           for await (const delta of generator) {
             if (delta !== void 0) {
               yield { version: errCount, delta };
@@ -7205,6 +7292,8 @@ Content: ${result.content}`
       kind: import_api22.SpanKind.SERVER,
       attributes
     });
+    const currentContext = import_api22.context.active();
+    const traceContext = import_api22.trace.setSpan(currentContext, span);
     try {
       if (!this.excludeContentFromTrace) {
         span.addEvent("input", { content: JSON.stringify(values, null, 2) });
@@ -7216,7 +7305,8 @@ Content: ${result.content}`
           ...options,
           functions
         },
-        span
+        span,
+        traceContext
       );
       if (!this.excludeContentFromTrace) {
         span.addEvent("output", {
@@ -8101,7 +8191,7 @@ var AxDBBase = class {
     if (!this.tracer) {
       return await this._upsert(req, update);
     }
-    return await this.tracer?.startActiveSpan(
+    return await this.tracer.startActiveSpan(
       "DB Upsert Request",
       {
         kind: import_api23.SpanKind.SERVER,
@@ -8135,7 +8225,7 @@ var AxDBBase = class {
     if (!this.tracer) {
       return await this._batchUpsert(req, update);
     }
-    return await this.tracer?.startActiveSpan(
+    return await this.tracer.startActiveSpan(
       "DB Batch Upsert Request",
       {
         kind: import_api23.SpanKind.SERVER,
@@ -8163,7 +8253,7 @@ var AxDBBase = class {
     if (!this.tracer) {
       return await this._query(req);
     }
-    return await this.tracer?.startActiveSpan(
+    return await this.tracer.startActiveSpan(
       "DB Query Request",
       {
         kind: import_api23.SpanKind.SERVER,
@@ -9140,24 +9230,24 @@ var AxJSInterpreter = class {
     this.permissions = permissions ?? [];
   }
   codeInterpreterJavascript(code) {
-    const context = { console };
+    const context3 = { console };
     if (this.permissions.includes("node:fs" /* FS */)) {
-      context.fs = _fs;
+      context3.fs = _fs;
     }
     if (this.permissions.includes("net" /* NET */)) {
-      context.http = _http;
-      context.https = _https;
+      context3.http = _http;
+      context3.https = _https;
     }
     if (this.permissions.includes("os" /* OS */)) {
-      context.os = _os;
+      context3.os = _os;
     }
     if (this.permissions.includes("crypto" /* CRYPTO */)) {
-      context.crypto = _crypto;
+      context3.crypto = _crypto;
     }
     if (this.permissions.includes("process" /* PROCESS */)) {
-      context.process = _process;
+      context3.process = _process;
     }
-    return (0, import_node_vm.runInNewContext)(`(function() { ${code} })()`, context);
+    return (0, import_node_vm.runInNewContext)(`(function() { ${code} })()`, context3);
   }
   toFunction() {
     return {
@@ -9940,9 +10030,9 @@ var colorLog6 = new ColorLog();
 var AxSimpleClassifierClass = class {
   name;
   context;
-  constructor(name, context) {
+  constructor(name, context3) {
     this.name = name;
-    this.context = context;
+    this.context = context3;
   }
   getName() {
     return this.name;
@@ -11946,20 +12036,20 @@ var AxRAG = class extends AxChainOfThought {
     this.register(this.genQuery);
   }
   async forward(ai, { question }, options) {
-    let context = [];
+    let context3 = [];
     for (let i = 0; i < this.maxHops; i++) {
       const { query } = await this.genQuery.forward(
         ai,
         {
-          context,
+          context: context3,
           question
         },
         options
       );
       const val = await this.queryFn(query);
-      context = AxStringUtil.dedup([...context, val]);
+      context3 = AxStringUtil.dedup([...context3, val]);
     }
-    return super.forward(ai, { context, question }, options);
+    return super.forward(ai, { context: context3, question }, options);
   }
 };
 // Annotate the CommonJS export names for ESM import in node: