@ax-llm/ax 11.0.42 → 11.0.44

package/index.js

@@ -5,6 +5,7 @@ import { SpanKind } from "@opentelemetry/api";
 var axSpanAttributes = {
   // LLM
   LLM_SYSTEM: "gen_ai.system",
+  LLM_OPERATION_NAME: "gen_ai.operation.name",
   LLM_REQUEST_MODEL: "gen_ai.request.model",
   LLM_REQUEST_MAX_TOKENS: "gen_ai.request.max_tokens",
   LLM_REQUEST_TEMPERATURE: "gen_ai.request.temperature",
@@ -14,8 +15,10 @@ var axSpanAttributes = {
   LLM_REQUEST_STOP_SEQUENCES: "gen_ai.request.stop_sequences",
   LLM_REQUEST_LLM_IS_STREAMING: "gen_ai.request.llm_is_streaming",
   LLM_REQUEST_TOP_P: "gen_ai.request.top_p",
-  LLM_USAGE_PROMPT_TOKENS: "gen_ai.usage.prompt_tokens",
-  LLM_USAGE_COMPLETION_TOKENS: "gen_ai.usage.completion_tokens",
+  LLM_USAGE_INPUT_TOKENS: "gen_ai.usage.input_tokens",
+  LLM_USAGE_OUTPUT_TOKENS: "gen_ai.usage.output_tokens",
+  LLM_USAGE_TOTAL_TOKENS: "gen_ai.usage.total_tokens",
+  LLM_USAGE_THOUGHTS_TOKENS: "gen_ai.usage.thoughts_tokens",
   // Vector DB
   DB_SYSTEM: "db.system",
   DB_TABLE: "db.table",
@@ -38,7 +41,13 @@ var axSpanAttributes = {
   DB_QUERY_RESULT_DOCUMENT: "db.query.result.document"
 };
 var axSpanEvents = {
-  LLM_PROMPT: "gen_ai.prompt"
+  GEN_AI_USER_MESSAGE: "gen_ai.user.message",
+  GEN_AI_SYSTEM_MESSAGE: "gen_ai.system.message",
+  GEN_AI_ASSISTANT_MESSAGE: "gen_ai.assistant.message",
+  GEN_AI_TOOL_MESSAGE: "gen_ai.tool.message",
+  // For tool messages in request & response tool calls
+  GEN_AI_CHOICE: "gen_ai.choice",
+  GEN_AI_USAGE: "gen_ai.usage"
 };
 var AxLLMRequestTypeValues = /* @__PURE__ */ ((AxLLMRequestTypeValues2) => {
   AxLLMRequestTypeValues2["COMPLETION"] = "completion";
@@ -482,9 +491,6 @@ var apiCall = async (api, json) => {
             } finally {
               clearTimeout(timeoutId);
               reader.releaseLock();
-              if (api.span?.isRecording()) {
-                api.span.end();
-              }
             }
           }
           read();
@@ -532,9 +538,6 @@ var apiCall = async (api, json) => {
       if (timeoutId !== void 0) {
         clearTimeout(timeoutId);
       }
-      if (api.span?.isRecording()) {
-        api.span.end();
-      }
     }
   }
 };
@@ -750,6 +753,7 @@ var AxBaseAI = class {
   fetch;
   tracer;
   timeout;
+  excludeContentFromTrace;
   models;
   modelInfo;
   modelUsage;
@@ -810,13 +814,16 @@ var AxBaseAI = class {
     this.fetch = options.fetch;
     this.timeout = options.timeout;
     this.tracer = options.tracer;
+    this.excludeContentFromTrace = options.excludeContentFromTrace;
   }
   getOptions() {
     return {
       debug: this.debug,
       rateLimiter: this.rt,
       fetch: this.fetch,
-      tracer: this.tracer
+      tracer: this.tracer,
+      timeout: this.timeout,
+      excludeContentFromTrace: this.excludeContentFromTrace
     };
   }
   getModelList() {
@@ -914,13 +921,18 @@ var AxBaseAI = class {
     if (!canStream) {
       modelConfig.stream = false;
     }
+    const canSetThinkingTokenBudget = this.getFeatures(model).thinkingTokenBudget;
+    if (!canSetThinkingTokenBudget && options?.thinkingTokenBudget) {
+      throw new Error("Thinking token budget is not supported for this model");
+    }
     if (this.tracer) {
       return await this.tracer?.startActiveSpan(
-        "Chat Request",
+        "AI Chat Request",
         {
           kind: SpanKind.SERVER,
           attributes: {
             [axSpanAttributes.LLM_SYSTEM]: this.name,
+            [axSpanAttributes.LLM_OPERATION_NAME]: "chat",
             [axSpanAttributes.LLM_REQUEST_MODEL]: model,
             [axSpanAttributes.LLM_REQUEST_MAX_TOKENS]: modelConfig.maxTokens,
             [axSpanAttributes.LLM_REQUEST_TEMPERATURE]: modelConfig.temperature,
@@ -930,17 +942,10 @@ var AxBaseAI = class {
             [axSpanAttributes.LLM_REQUEST_PRESENCE_PENALTY]: modelConfig.presencePenalty,
             [axSpanAttributes.LLM_REQUEST_STOP_SEQUENCES]: modelConfig.stopSequences?.join(", "),
             [axSpanAttributes.LLM_REQUEST_LLM_IS_STREAMING]: modelConfig.stream
-            // [AxSpanAttributes.LLM_PROMPTS]: _req.chatPrompt
-            //   ?.map((v) => v.content)
-            //   .join('\n')
           }
         },
         async (span) => {
-          try {
-            return await this._chat2(model, modelConfig, req, options, span);
-          } finally {
-            span.end();
-          }
+          return await this._chat2(model, modelConfig, req, options, span);
         }
       );
     }
@@ -986,6 +991,9 @@ var AxBaseAI = class {
         req,
         options
       );
+      if (span?.isRecording()) {
+        setChatRequestEvents(chatReq, span, this.excludeContentFromTrace);
+      }
       const res2 = await apiCall(
         {
           name: apiConfig.name,
@@ -1023,7 +1031,7 @@ var AxBaseAI = class {
         }
         this.modelUsage = res2.modelUsage;
         if (span?.isRecording()) {
-          setResponseAttr(res2, span);
+          setChatResponseEvents(res2, span, this.excludeContentFromTrace);
         }
         if (debug) {
           logResponse(res2);
@@ -1034,6 +1042,9 @@ var AxBaseAI = class {
         if (debug) {
           process.stdout.write("\n");
         }
+        if (span?.isRecording()) {
+          span.end();
+        }
       };
       const st = rv.pipeThrough(
         new RespTransformStream(
@@ -1048,16 +1059,26 @@ var AxBaseAI = class {
     }
     const res = this.aiImpl.createChatResp(rv);
     res.sessionId = options?.sessionId;
+    if (!res.modelUsage) {
+      const tokenUsage = this.aiImpl.getTokenUsage();
+      if (tokenUsage) {
+        res.modelUsage = {
+          ai: this.name,
+          model,
+          tokens: tokenUsage
+        };
+      }
+    }
     if (res.modelUsage) {
       this.modelUsage = res.modelUsage;
     }
     if (span?.isRecording()) {
-      setResponseAttr(res, span);
+      setChatResponseEvents(res, span, this.excludeContentFromTrace);
+      span.end();
     }
     if (debug) {
       logResponse(res);
     }
-    span?.end();
     return res;
   }
   async embed(req, options) {
@@ -1081,11 +1102,12 @@ var AxBaseAI = class {
     }
     if (this.tracer) {
       await this.tracer?.startActiveSpan(
-        "Embed Request",
+        "AI Embed Request",
         {
           kind: SpanKind.SERVER,
           attributes: {
             [axSpanAttributes.LLM_SYSTEM]: this.name,
+            [axSpanAttributes.LLM_OPERATION_NAME]: "embeddings",
             [axSpanAttributes.LLM_REQUEST_MODEL]: embedModel
           }
         },
@@ -1140,8 +1162,12 @@ var AxBaseAI = class {
       };
     }
     this.embedModelUsage = res.modelUsage;
-    if (span?.isRecording()) {
-      setResponseAttr(res, span);
+    if (span?.isRecording() && res.modelUsage?.tokens) {
+      span.addEvent(axSpanEvents.GEN_AI_USAGE, {
+        [axSpanAttributes.LLM_USAGE_INPUT_TOKENS]: res.modelUsage.tokens.promptTokens,
+        [axSpanAttributes.LLM_USAGE_OUTPUT_TOKENS]: res.modelUsage.tokens.completionTokens ?? 0,
+        [axSpanAttributes.LLM_USAGE_TOTAL_TOKENS]: res.modelUsage.tokens.totalTokens
+      });
     }
     span?.end();
     return res;
@@ -1165,11 +1191,113 @@ var AxBaseAI = class {
     return item && "embedModel" in item ? item.embedModel : void 0;
   }
 };
-function setResponseAttr(res, span) {
-  if (res.modelUsage) {
-    span.setAttributes({
-      [axSpanAttributes.LLM_USAGE_COMPLETION_TOKENS]: res.modelUsage.tokens?.completionTokens ?? 0,
-      [axSpanAttributes.LLM_USAGE_PROMPT_TOKENS]: res.modelUsage.tokens?.promptTokens
+function setChatRequestEvents(req, span, excludeContentFromTrace) {
+  const userMessages = [];
+  if (req.chatPrompt && Array.isArray(req.chatPrompt) && req.chatPrompt.length > 0) {
+    for (const prompt of req.chatPrompt) {
+      switch (prompt.role) {
+        case "system":
+          if (prompt.content) {
+            const eventData2 = {};
+            if (!excludeContentFromTrace) {
+              eventData2.content = prompt.content;
+            }
+            span.addEvent(axSpanEvents.GEN_AI_SYSTEM_MESSAGE, eventData2);
+          }
+          break;
+        case "user":
+          if (typeof prompt.content === "string") {
+            userMessages.push(prompt.content);
+          } else if (Array.isArray(prompt.content)) {
+            for (const part of prompt.content) {
+              if (part.type === "text") {
+                userMessages.push(part.text);
+              }
+            }
+          }
+          break;
+        case "assistant":
+          const functionCalls = prompt.functionCalls?.map((call) => {
+            return {
+              id: call.id,
+              type: call.type,
+              function: call.function.name,
+              arguments: call.function.params
+            };
+          });
+          if (functionCalls && functionCalls.length > 0) {
+            const eventData2 = {
+              function_calls: JSON.stringify(functionCalls, null, 2)
+            };
+            if (!excludeContentFromTrace && prompt.content) {
+              eventData2.content = prompt.content;
+            }
+            span.addEvent(axSpanEvents.GEN_AI_ASSISTANT_MESSAGE, eventData2);
+          } else if (prompt.content) {
+            const eventData2 = {};
+            if (!excludeContentFromTrace) {
+              eventData2.content = prompt.content;
+            }
+            span.addEvent(axSpanEvents.GEN_AI_ASSISTANT_MESSAGE, eventData2);
+          }
+          break;
+        case "function":
+          const eventData = {
+            id: prompt.functionId
+          };
+          if (!excludeContentFromTrace) {
+            eventData.content = prompt.result;
+          }
+          span.addEvent(axSpanEvents.GEN_AI_TOOL_MESSAGE, eventData);
+          break;
+      }
+    }
+  }
+  const userEventData = {};
+  if (!excludeContentFromTrace) {
+    userEventData.content = userMessages.join("\n");
+  }
+  span.addEvent(axSpanEvents.GEN_AI_USER_MESSAGE, userEventData);
+}
+function setChatResponseEvents(res, span, excludeContentFromTrace) {
+  if (res.modelUsage?.tokens) {
+    const thoughTokens = res.modelUsage.tokens.thoughtsTokens ? {
+      [axSpanAttributes.LLM_USAGE_THOUGHTS_TOKENS]: res.modelUsage.tokens.thoughtsTokens
+    } : {};
+    span.addEvent(axSpanEvents.GEN_AI_USAGE, {
+      [axSpanAttributes.LLM_USAGE_INPUT_TOKENS]: res.modelUsage.tokens.promptTokens,
+      [axSpanAttributes.LLM_USAGE_OUTPUT_TOKENS]: res.modelUsage.tokens.completionTokens ?? 0,
+      [axSpanAttributes.LLM_USAGE_TOTAL_TOKENS]: res.modelUsage.tokens.totalTokens,
+      ...thoughTokens
+    });
+  }
+  if (!res.results) {
+    return;
+  }
+  for (const [index, result] of res.results.entries()) {
+    const toolCalls = result.functionCalls?.map((call) => {
+      return {
+        id: call.id,
+        type: call.type,
+        function: call.function.name,
+        arguments: call.function.params
+      };
+    });
+    let message = {};
+    if (toolCalls && toolCalls.length > 0) {
+      if (!excludeContentFromTrace) {
+        message.content = result.content;
+      }
+      message.tool_calls = toolCalls;
+    } else {
+      if (!excludeContentFromTrace) {
+        message.content = result.content ?? "";
+      }
+    }
+    span.addEvent(axSpanEvents.GEN_AI_CHOICE, {
+      finish_reason: result.finishReason,
+      index,
+      message: JSON.stringify(message, null, 2)
     });
   }
 }
@@ -2562,7 +2690,7 @@ var AxAIDeepSeek = class extends AxAIOpenAIBase {
 
 // ai/google-gemini/types.ts
 var AxAIGoogleGeminiModel = /* @__PURE__ */ ((AxAIGoogleGeminiModel2) => {
-  AxAIGoogleGeminiModel2["Gemini25Pro"] = "gemini-2.5-pro-preview-03-25";
+  AxAIGoogleGeminiModel2["Gemini25Pro"] = "gemini-2.5-pro-preview-05-06";
   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";
@@ -2611,7 +2739,7 @@ var AxAIGoogleGeminiEmbedTypes = /* @__PURE__ */ ((AxAIGoogleGeminiEmbedTypes2)
 // ai/google-gemini/info.ts
 var axModelInfoGoogleGemini = [
   {
-    name: "gemini-2.5-pro-preview-03-25" /* Gemini25Pro */,
+    name: "gemini-2.5-pro-preview-05-06" /* Gemini25Pro */,
     currency: "usd",
     characterIsToken: false,
     promptTokenCostPer1M: 2.5,
@@ -2695,7 +2823,7 @@ var safetySettings = [
   }
 ];
 var axAIGoogleGeminiDefaultConfig = () => structuredClone({
-  model: "gemini-2.0-flash" /* Gemini20Flash */,
+  model: "gemini-2.5-flash-preview-04-17" /* Gemini25Flash */,
   embedModel: "text-embedding-005" /* TextEmbedding005 */,
   safetySettings,
   ...axBaseAIDefaultConfig()
@@ -2736,7 +2864,7 @@ var AxAIGoogleGeminiImpl = class {
       n: config.n
     };
   }
-  createChatReq = (req) => {
+  createChatReq = (req, config) => {
     const model = req.model;
     const stream = req.modelConfig?.stream ?? this.config.stream;
     if (!req.chatPrompt || req.chatPrompt.length === 0) {
@@ -2847,6 +2975,9 @@ var AxAIGoogleGeminiImpl = class {
         }
       });
     }
+    if (this.options?.urlContext) {
+      tools.push({ url_context: {} });
+    }
     if (tools.length === 0) {
       tools = void 0;
     }
@@ -2872,6 +3003,16 @@ var AxAIGoogleGeminiImpl = class {
     } else if (tools && tools.length > 0) {
       toolConfig = { function_calling_config: { mode: "AUTO" } };
     }
+    const thinkingConfig = {};
+    if (this.config.thinking?.includeThoughts) {
+      thinkingConfig.includeThoughts = true;
+    }
+    if (this.config.thinking?.thinkingTokenBudget) {
+      thinkingConfig.thinkingBudget = this.config.thinking.thinkingTokenBudget;
+    }
+    if (config.thinkingTokenBudget) {
+      thinkingConfig.thinkingBudget = config.thinkingTokenBudget;
+    }
     const generationConfig = {
       maxOutputTokens: req.modelConfig?.maxTokens ?? this.config.maxTokens,
       temperature: req.modelConfig?.temperature ?? this.config.temperature,
@@ -2881,9 +3022,7 @@ var AxAIGoogleGeminiImpl = class {
       candidateCount: 1,
       stopSequences: req.modelConfig?.stopSequences ?? this.config.stopSequences,
       responseMimeType: "text/plain",
-      ...this.config.thinkingConfig && {
-        thinkingConfig: this.config.thinkingConfig
-      }
+      ...thinkingConfig ? { thinkingConfig } : {}
     };
     const safetySettings2 = this.config.safetySettings;
     const reqValue = {
@@ -2964,7 +3103,11 @@ var AxAIGoogleGeminiImpl = class {
         }
         for (const part of candidate.content.parts) {
           if ("text" in part) {
-            result.content = part.text;
+            if ("thought" in part && part.thought) {
+              result.thought = part.text;
+            } else {
+              result.content = part.text;
+            }
             continue;
           }
           if ("functionCall" in part) {
@@ -2987,7 +3130,8 @@ var AxAIGoogleGeminiImpl = class {
       this.tokensUsed = {
         totalTokens: resp.usageMetadata.totalTokenCount,
         promptTokens: resp.usageMetadata.promptTokenCount,
-        completionTokens: resp.usageMetadata.candidatesTokenCount
+        completionTokens: resp.usageMetadata.candidatesTokenCount,
+        thoughtsTokens: resp.usageMetadata.thoughtsTokenCount
       };
     }
     return { results };
@@ -4578,6 +4722,14 @@ var AxSignature = class _AxSignature {
   };
   hash = () => this.sigHash;
   toString = () => this.sigString;
+  toJSON = () => {
+    return {
+      id: this.hash(),
+      description: this.description,
+      inputFields: this.inputFields,
+      outputFields: this.outputFields
+    };
+  };
 };
 function renderField(field) {
   let result = field.name;
@@ -6118,7 +6270,7 @@ var parseFunctions = (newFuncs, existingFuncs) => {
   }
   return [...existingFuncs ?? [], ...functions];
 };
-var processFunctions = async (ai, functionList, functionCalls, mem, sessionId, traceId) => {
+var processFunctions = async (ai, functionList, functionCalls, mem, sessionId, traceId, span, excludeContentFromTelemetry) => {
   const funcProc = new AxFunctionProcessor(functionList);
   const functionsExecuted = /* @__PURE__ */ new Set();
   const promises = functionCalls.map((func) => {
@@ -6127,6 +6279,16 @@ var processFunctions = async (ai, functionList, functionCalls, mem, sessionId, t
     }
     const promise = funcProc.execute(func, { sessionId, traceId, ai }).then((functionResult) => {
       functionsExecuted.add(func.name.toLowerCase());
+      if (span) {
+        const eventData = {
+          name: func.name
+        };
+        if (!excludeContentFromTelemetry) {
+          eventData.args = func.args;
+          eventData.result = functionResult ?? "";
+        }
+        span.addEvent("function.call", eventData);
+      }
       return {
         role: "function",
         result: functionResult ?? "",
@@ -6135,6 +6297,17 @@ var processFunctions = async (ai, functionList, functionCalls, mem, sessionId, t
     }).catch((e) => {
       if (e instanceof FunctionError) {
         const result = e.getFixingInstructions();
+        if (span) {
+          const errorEventData = {
+            name: func.name,
+            message: e.toString()
+          };
+          if (!excludeContentFromTelemetry) {
+            errorEventData.args = func.args;
+            errorEventData.fixing_instructions = result;
+          }
+          span.addEvent("function.error", errorEventData);
+        }
         mem.add(
           {
             role: "function",
@@ -6192,6 +6365,8 @@ var AxGen = class extends AxProgramWithSignature {
   functionsExecuted = /* @__PURE__ */ new Set();
   fieldProcessors = [];
   streamingFieldProcessors = [];
+  values = {};
+  excludeContentFromTrace = false;
   constructor(signature, options) {
     super(signature, { description: options?.description });
     this.options = options;
@@ -6201,6 +6376,7 @@ var AxGen = class extends AxProgramWithSignature {
     );
     this.asserts = this.options?.asserts ?? [];
     this.streamingAsserts = this.options?.streamingAsserts ?? [];
+    this.excludeContentFromTrace = options?.excludeContentFromTrace ?? false;
     this.usage = [];
     if (options?.functions) {
       this.functions = parseFunctions(options.functions);
@@ -6249,7 +6425,8 @@ var AxGen = class extends AxProgramWithSignature {
       rateLimiter,
       stream,
       functions: _functions,
-      functionCall: _functionCall
+      functionCall: _functionCall,
+      thinkingTokenBudget
     } = options ?? {};
     const chatPrompt = mem?.history(sessionId) ?? [];
     if (chatPrompt.length === 0) {
@@ -6270,7 +6447,8 @@ var AxGen = class extends AxProgramWithSignature {
         traceId,
         rateLimiter,
         stream,
-        debug: false
+        debug: false,
+        thinkingTokenBudget
       }
     );
     return res;
@@ -6278,7 +6456,8 @@ var AxGen = class extends AxProgramWithSignature {
   async *forwardCore({
     ai,
     mem,
-    options
+    options,
+    span
   }) {
     const { sessionId, traceId, functions: _functions } = options ?? {};
     const fastFail = options?.fastFail ?? this.options?.fastFail;
@@ -6298,7 +6477,8 @@ var AxGen = class extends AxProgramWithSignature {
         traceId,
         sessionId,
         functions,
-        fastFail
+        fastFail,
+        span
       });
     } else {
       yield await this.processResponse({
@@ -6308,7 +6488,8 @@ var AxGen = class extends AxProgramWithSignature {
         mem,
         traceId,
         sessionId,
-        functions
+        functions,
+        span
       });
     }
   }
@@ -6320,11 +6501,12 @@ var AxGen = class extends AxProgramWithSignature {
     sessionId,
     traceId,
     functions,
-    fastFail
+    fastFail,
+    span
   }) {
     const streamingValidation = fastFail ?? ai.getFeatures(model).functionCot !== true;
     const functionCalls = [];
-    const values = {};
+    this.values = {};
     const xstate = {
       extractedFields: [],
       streamedIndex: {},
@@ -6351,6 +6533,9 @@ var AxGen = class extends AxProgramWithSignature {
           sessionId
         );
       } else if (result.content) {
+        if (result.thought && result.thought.length > 0) {
+          yield { thought: result.thought };
+        }
         content += result.content;
         mem.updateResult(
           { name: result.name, content, delta: result.content },
@@ -6358,7 +6543,7 @@ var AxGen = class extends AxProgramWithSignature {
         );
         const skip = streamingExtractValues(
           this.signature,
-          values,
+          this.values,
           xstate,
           content,
           streamingValidation
@@ -6379,12 +6564,20 @@ var AxGen = class extends AxProgramWithSignature {
             content,
             xstate,
             mem,
-            values,
+            this.values,
             sessionId
           );
         }
-        yield* streamValues(this.signature, content, values, xstate);
-        await assertAssertions(this.asserts, values);
+        yield* streamValues(
+          this.signature,
+          content,
+          this.values,
+          xstate
+        );
+        await assertAssertions(this.asserts, this.values);
+      } else if (result.thought && result.thought.length > 0) {
+        this.values.thought = this.values.thought ?? "" + result.thought;
+        yield { thought: result.thought };
       }
       if (result.finishReason === "length") {
         throw new Error(
@@ -6393,7 +6586,7 @@ Content: ${content}`
         );
       }
     }
-    const funcs = parseFunctionCalls(ai, functionCalls, values, model);
+    const funcs = parseFunctionCalls(ai, functionCalls, this.values, model);
     if (funcs) {
       if (!functions) {
         throw new Error("Functions are not defined");
@@ -6404,22 +6597,24 @@ Content: ${content}`
         funcs,
         mem,
         sessionId,
-        traceId
+        traceId,
+        span,
+        this.excludeContentFromTrace
       );
       this.functionsExecuted = /* @__PURE__ */ new Set([...this.functionsExecuted, ...fx]);
     } else {
-      streamingExtractFinalValue(this.signature, values, xstate, content);
+      streamingExtractFinalValue(this.signature, this.values, xstate, content);
       await assertStreamingAssertions(
         this.streamingAsserts,
         xstate,
         content,
         true
       );
-      await assertAssertions(this.asserts, values);
+      await assertAssertions(this.asserts, this.values);
       if (this.fieldProcessors.length) {
         await processFieldProcessors(
           this.fieldProcessors,
-          values,
+          this.values,
           mem,
           sessionId
         );
@@ -6430,12 +6625,17 @@ Content: ${content}`
           content,
           xstate,
           mem,
-          values,
+          this.values,
           sessionId,
           true
         );
       }
-      yield* streamValues(this.signature, content, values, xstate);
+      yield* streamValues(
+        this.signature,
+        content,
+        this.values,
+        xstate
+      );
     }
   }
   async processResponse({
@@ -6444,9 +6644,10 @@ Content: ${content}`
     mem,
     sessionId,
     traceId,
-    functions
+    functions,
+    span
   }) {
-    const values = {};
+    this.values = {};
     let results = res.results ?? [];
     if (results.length > 1) {
       results = results.filter((r) => r.functionCalls);
@@ -6457,7 +6658,7 @@ Content: ${content}`
       }
       mem.addResult(result, sessionId);
       if (result.functionCalls?.length) {
-        const funcs = parseFunctionCalls(ai, result.functionCalls, values);
+        const funcs = parseFunctionCalls(ai, result.functionCalls, this.values);
         if (funcs) {
           if (!functions) {
             throw new Error("Functions are not defined");
@@ -6468,17 +6669,22 @@ Content: ${content}`
             funcs,
             mem,
             sessionId,
-            traceId
+            traceId,
+            span,
+            this.excludeContentFromTrace
           );
           this.functionsExecuted = /* @__PURE__ */ new Set([...this.functionsExecuted, ...fx]);
         }
       } else if (result.content) {
-        extractValues(this.signature, values, result.content);
-        await assertAssertions(this.asserts, values);
+        if (result.thought && result.thought.length > 0) {
+          this.values.thought = result.thought;
+        }
+        extractValues(this.signature, this.values, result.content);
+        await assertAssertions(this.asserts, this.values);
         if (this.fieldProcessors.length) {
           await processFieldProcessors(
             this.fieldProcessors,
-            values,
+            this.values,
             mem,
             sessionId
           );
@@ -6491,13 +6697,12 @@ Content: ${result.content}`
         );
       }
     }
-    const publicValues = { ...values };
     for (const field of this.signature.getOutputFields()) {
       if (field.isInternal) {
-        delete publicValues[field.name];
+        delete this.values[field.name];
       }
     }
-    return { ...values };
+    return { ...this.values };
   }
   async *_forward2(ai, values, options, span) {
     const stopFunction = (options?.stopFunction ?? this.options?.stopFunction)?.toLowerCase();
@@ -6523,7 +6728,7 @@ Content: ${result.content}`
     multiStepLoop: for (let n = 0; n < maxSteps; n++) {
       for (let errCount = 0; errCount < maxRetries; errCount++) {
         try {
-          const generator = this.forwardCore({ options, ai, mem });
+          const generator = this.forwardCore({ options, ai, mem, span });
           for await (const delta of generator) {
             if (delta !== void 0) {
               yield { version: errCount, delta };
@@ -6547,10 +6752,22 @@ Content: ${result.content}`
           if (e instanceof ValidationError) {
             errorFields = e.getFixingInstructions();
             err = e;
+            if (span) {
+              span.addEvent("validation.error", {
+                message: e.toString(),
+                fixing_instructions: errorFields?.map((f) => f.title).join(", ") ?? ""
+              });
+            }
           } else if (e instanceof AxAssertionError) {
             const e1 = e;
             errorFields = e1.getFixingInstructions();
             err = e;
+            if (span) {
+              span.addEvent("assertion.error", {
+                message: e1.toString(),
+                fixing_instructions: errorFields?.map((f) => f.title).join(", ") ?? ""
+              });
+            }
           } else if (e instanceof AxAIServiceStreamTerminatedError) {
           } else {
             throw enhanceError(e, ai, this.signature);
@@ -6605,14 +6822,25 @@ Content: ${result.content}`
     }
     const funcNames = functions?.map((f) => f.name).join(",");
     const attributes = {
-      "generate.signature": this.signature.toString(),
-      "generate.functions": funcNames ?? ""
+      signature: JSON.stringify(this.signature.toJSON(), null, 2),
+      ...this.examples ? { examples: JSON.stringify(this.examples, null, 2) } : {},
+      ...funcNames ? { provided_functions: funcNames } : {},
+      ...options?.model ? { model: options.model } : {},
+      ...options?.thinkingTokenBudget ? { thinking_token_budget: options.thinkingTokenBudget } : {},
+      ...options?.maxSteps ? { max_steps: options.maxSteps } : {},
+      ...options?.maxRetries ? { max_retries: options.maxRetries } : {},
+      ...options?.fastFail ? { fast_fail: options.fastFail } : {}
     };
-    const span = tracer.startSpan("Generate", {
+    const traceLabel = options.traceLabel ?? this.options?.traceLabel;
+    const spanName = traceLabel ? `${traceLabel} (AxGen)` : "AxGen";
+    const span = tracer.startSpan(spanName, {
       kind: SpanKind2.SERVER,
       attributes
     });
     try {
+      if (!this.excludeContentFromTrace) {
+        span.addEvent("input", { content: JSON.stringify(values, null, 2) });
+      }
       yield* this._forward2(
         ai,
         values,
@@ -6622,6 +6850,11 @@ Content: ${result.content}`
         },
         span
       );
+      if (!this.excludeContentFromTrace) {
+        span.addEvent("output", {
+          content: JSON.stringify(this.values, null, 2)
+        });
+      }
     } finally {
       span.end();
     }