@ax-llm/ax 11.0.47 → 11.0.48

package/index.d.cts

@@ -1,5 +1,5 @@
 import { ReadableStream as ReadableStream$1 } from 'node:stream/web';
-import { Span, Tracer } from '@opentelemetry/api';
+import { Span, Context, Tracer } from '@opentelemetry/api';
 import { ReadableStream } from 'stream/web';
 
 interface RetryConfig {
@@ -332,7 +332,7 @@ type AxRateLimiterFunction = <T = unknown>(reqFunc: () => Promise<T | ReadableSt
 }>) => Promise<T | ReadableStream$1<T>>;
 type AxAIPromptConfig = {
     stream?: boolean;
-    thinkingTokenBudget?: 'minimal' | 'low' | 'medium' | 'high';
+    thinkingTokenBudget?: 'minimal' | 'low' | 'medium' | 'high' | 'highest';
 };
 type AxAIServiceOptions = {
     debug?: boolean;
@@ -351,6 +351,7 @@ type AxAIServiceActionOptions<TModel = unknown, TEmbedModel = unknown> = {
     debug?: boolean;
     debugHideSystemPrompt?: boolean;
     hideThought?: boolean;
+    traceContext?: Context;
 };
 interface AxAIService<TModel = unknown, TEmbedModel = unknown> {
     getId(): string;
@@ -1691,7 +1692,7 @@ type AxProgramForwardOptions = {
     fastFail?: boolean;
     debug?: boolean;
     debugHideSystemPrompt?: boolean;
-    thinkingTokenBudget?: 'minimal' | 'low' | 'medium' | 'high';
+    thinkingTokenBudget?: 'minimal' | 'low' | 'medium' | 'high' | 'highest';
     traceLabel?: string;
 };
 type AxProgramStreamingForwardOptions = Omit<AxProgramForwardOptions, 'stream'>;

package/index.d.ts

@@ -1,5 +1,5 @@
 import { ReadableStream as ReadableStream$1 } from 'node:stream/web';
-import { Span, Tracer } from '@opentelemetry/api';
+import { Span, Context, Tracer } from '@opentelemetry/api';
 import { ReadableStream } from 'stream/web';
 
 interface RetryConfig {
@@ -332,7 +332,7 @@ type AxRateLimiterFunction = <T = unknown>(reqFunc: () => Promise<T | ReadableSt
 }>) => Promise<T | ReadableStream$1<T>>;
 type AxAIPromptConfig = {
     stream?: boolean;
-    thinkingTokenBudget?: 'minimal' | 'low' | 'medium' | 'high';
+    thinkingTokenBudget?: 'minimal' | 'low' | 'medium' | 'high' | 'highest';
 };
 type AxAIServiceOptions = {
     debug?: boolean;
@@ -351,6 +351,7 @@ type AxAIServiceActionOptions<TModel = unknown, TEmbedModel = unknown> = {
     debug?: boolean;
     debugHideSystemPrompt?: boolean;
     hideThought?: boolean;
+    traceContext?: Context;
 };
 interface AxAIService<TModel = unknown, TEmbedModel = unknown> {
     getId(): string;
@@ -1691,7 +1692,7 @@ type AxProgramForwardOptions = {
     fastFail?: boolean;
     debug?: boolean;
     debugHideSystemPrompt?: boolean;
-    thinkingTokenBudget?: 'minimal' | 'low' | 'medium' | 'high';
+    thinkingTokenBudget?: 'minimal' | 'low' | 'medium' | 'high' | 'highest';
     traceLabel?: string;
 };
 type AxProgramStreamingForwardOptions = Omit<AxProgramForwardOptions, 'stream'>;

package/index.js

@@ -1,5 +1,5 @@
 // ai/base.ts
-import { SpanKind } from "@opentelemetry/api";
+import { context, SpanKind } from "@opentelemetry/api";
 
 // trace/trace.ts
 var axSpanAttributes = {
@@ -205,14 +205,14 @@ var defaultRetryConfig = {
 var defaultTimeoutMs = 3e4;
 var textDecoderStream = TextDecoderStreamNative ?? 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;
@@ -234,11 +234,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;
@@ -246,11 +246,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;
@@ -258,33 +258,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;
   }
 };
@@ -330,14 +330,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();
@@ -364,16 +362,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;
       }
@@ -399,12 +395,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) {
@@ -424,15 +418,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;
@@ -517,16 +507,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;
       }
@@ -927,7 +915,7 @@ var AxBaseAI = class {
       modelConfig.stream = false;
     }
     if (this.tracer) {
-      return await this.tracer?.startActiveSpan(
+      return await this.tracer.startActiveSpan(
         "AI Chat Request",
         {
           kind: SpanKind.SERVER,
@@ -945,6 +933,7 @@ var AxBaseAI = class {
             [axSpanAttributes.LLM_REQUEST_LLM_IS_STREAMING]: modelConfig.stream
           }
         },
+        options?.traceContext ?? context.active(),
         async (span) => {
           return await this._chat2(model, modelConfig, req, options, span);
         }
@@ -1126,6 +1115,7 @@ var AxBaseAI = class {
             [axSpanAttributes.LLM_REQUEST_MODEL]: embedModel
           }
         },
+        options?.traceContext ?? context.active(),
         async (span) => {
           try {
             return await this._embed2(embedModel, req, options, span);
@@ -2076,6 +2066,10 @@ var AxAIOpenAIImpl = class {
           break;
         case "high":
           reqValue.reasoning_effort = "high";
+          break;
+        case "highest":
+          reqValue.reasoning_effort = "high";
+          break;
       }
     }
     if (this.chatReqUpdater) {
@@ -3114,16 +3108,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;
       }
     }
@@ -4189,7 +4186,11 @@ var AxAIGrok = class extends AxAIOpenAIBase {
 
 // dsp/generate.ts
 import { ReadableStream as ReadableStream3 } from "node:stream/web";
-import { SpanKind as SpanKind2 } from "@opentelemetry/api";
+import {
+  context as context2,
+  SpanKind as SpanKind2,
+  trace
+} from "@opentelemetry/api";
 
 // ai/util.ts
 function mergeFunctionCalls(functionCalls, functionCallDeltas) {
@@ -4524,9 +4525,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() {
@@ -6625,7 +6626,8 @@ var AxGen = class extends AxProgramWithSignature {
   async forwardSendRequest({
     ai,
     mem,
-    options
+    options,
+    traceContext
   }) {
     const {
       sessionId,
@@ -6658,7 +6660,8 @@ var AxGen = class extends AxProgramWithSignature {
         rateLimiter,
         stream,
         debug: false,
-        thinkingTokenBudget
+        thinkingTokenBudget,
+        traceContext
       }
     );
     return res;
@@ -6667,7 +6670,8 @@ var AxGen = class extends AxProgramWithSignature {
     ai,
     mem,
     options,
-    span
+    span,
+    traceContext
   }) {
     const { sessionId, traceId, functions: _functions } = options ?? {};
     const fastFail = options?.fastFail ?? this.options?.fastFail;
@@ -6676,7 +6680,8 @@ var AxGen = class extends AxProgramWithSignature {
     const res = await this.forwardSendRequest({
       ai,
       mem,
-      options
+      options,
+      traceContext
     });
     if (res instanceof ReadableStream3) {
       yield* this.processStreamingResponse({
@@ -6914,7 +6919,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;
@@ -6938,7 +6943,13 @@ 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, span });
+          const generator = this.forwardCore({
+            options,
+            ai,
+            mem,
+            span,
+            traceContext
+          });
           for await (const delta of generator) {
             if (delta !== void 0) {
               yield { version: errCount, delta };
@@ -7047,6 +7058,8 @@ Content: ${result.content}`
       kind: SpanKind2.SERVER,
       attributes
     });
+    const currentContext = context2.active();
+    const traceContext = trace.setSpan(currentContext, span);
     try {
       if (!this.excludeContentFromTrace) {
         span.addEvent("input", { content: JSON.stringify(values, null, 2) });
@@ -7058,7 +7071,8 @@ Content: ${result.content}`
           ...options,
           functions
         },
-        span
+        span,
+        traceContext
       );
       if (!this.excludeContentFromTrace) {
         span.addEvent("output", {
@@ -7943,7 +7957,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: SpanKind3.SERVER,
@@ -7977,7 +7991,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: SpanKind3.SERVER,
@@ -8005,7 +8019,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: SpanKind3.SERVER,
@@ -8982,24 +8996,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 runInNewContext(`(function() { ${code} })()`, context);
+    return runInNewContext(`(function() { ${code} })()`, context3);
   }
   toFunction() {
     return {
@@ -9782,9 +9796,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;
@@ -11788,20 +11802,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);
   }
 };
 export {