@ax-llm/ax 10.0.45 → 10.0.47

package/index.js

@@ -74,16 +74,16 @@ var AxSpanKindValues = /* @__PURE__ */ ((AxSpanKindValues2) => {
 })(AxSpanKindValues || {});
 
 // util/apicall.ts
-import path from "path";
 import {
   ReadableStream,
-  TextDecoderStream as TextDecoderStreamNative
+  TextDecoderStream as TextDecoderStreamNative,
+  TransformStream as TransformStream3
 } from "stream/web";
 import "@opentelemetry/api";
 
 // util/sse.ts
-import { TransformStream as TransformStream2 } from "stream/web";
-var SSEParser = class extends TransformStream2 {
+import { TransformStream } from "stream/web";
+var SSEParser = class extends TransformStream {
   buffer = "";
   currentEvent = { rawData: "" };
   dataParser;
@@ -173,7 +173,7 @@ var SSEParser = class extends TransformStream2 {
 
 // util/stream.ts
 import {
-  TransformStream as TransformStream3
+  TransformStream as TransformStream2
 } from "stream/web";
 var TextDecodeTransformer = class {
   decoder;
@@ -196,7 +196,7 @@ var TextDecodeTransformer = class {
     }
   }
 };
-var TextDecoderStreamPolyfill = class extends TransformStream3 {
+var TextDecoderStreamPolyfill = class extends TransformStream2 {
   constructor() {
     super(new TextDecodeTransformer());
   }
@@ -207,7 +207,6 @@ var defaultRetryConfig = {
   maxRetries: 3,
   initialDelayMs: 1e3,
   maxDelayMs: 6e4,
-  // Increased to 60 seconds
   backoffFactor: 2,
   retryableStatusCodes: [500, 408, 429, 502, 503, 504]
 };
@@ -218,18 +217,19 @@ var AxAIServiceError = class extends Error {
     super(message);
     this.url = url;
     this.requestBody = requestBody;
-    this.context = context;
     this.name = "AxAIServiceError";
     this.timestamp = (/* @__PURE__ */ new Date()).toISOString();
     this.errorId = crypto.randomUUID();
+    this.context = context;
   }
   timestamp;
   errorId;
+  context;
   toString() {
     return `${this.name} [${this.errorId}]: ${this.message}
 Timestamp: ${this.timestamp}
 URL: ${this.url}${this.requestBody ? `
-Request Body: ${JSON.stringify(this.requestBody, null, 2)}` : ""}${this.context ? `
+Request Body: ${JSON.stringify(this.requestBody, null, 2)}` : ""}${Object.keys(this.context).length ? `
 Context: ${JSON.stringify(this.context, null, 2)}` : ""}`;
   }
   toJSON() {
@@ -317,14 +317,31 @@ function updateRetryMetrics(metrics) {
   metrics.retryCount++;
   metrics.lastRetryTime = Date.now();
 }
+function shouldRetry(error, status, attempt, config) {
+  if (attempt >= config.maxRetries) return false;
+  if (status && config.retryableStatusCodes.includes(status)) return true;
+  return error instanceof AxAIServiceNetworkError && !(error instanceof AxAIServiceAuthenticationError);
+}
 var apiCall = async (api, json) => {
   const retryConfig = { ...defaultRetryConfig, ...api.retry };
   const timeoutMs = api.timeout ?? defaultTimeoutMs;
   const metrics = createRequestMetrics();
+  let timeoutId;
   const baseUrl = new URL(process.env["PROXY"] ?? api.url);
-  const apiPath = path.join(baseUrl.pathname, api.name ?? "/", baseUrl.search);
+  const apiPath = [baseUrl.pathname, api.name].filter(Boolean).join("/").replace(/\/+/g, "/");
   const apiUrl = new URL(apiPath, baseUrl);
   const requestId = crypto.randomUUID();
+  if (api.validateRequest) {
+    const isValid = await api.validateRequest(json);
+    if (!isValid) {
+      throw new AxAIServiceResponseError(
+        "Invalid request data",
+        apiUrl.href,
+        json,
+        { validation: "request" }
+      );
+    }
+  }
   if (api.span?.isRecording()) {
     api.span.setAttributes({
       "http.request.method": api.put ? "PUT" : "POST",
@@ -336,7 +353,7 @@ var apiCall = async (api, json) => {
   let attempt = 0;
   while (true) {
     const controller = new AbortController();
-    let timeoutId = setTimeout(() => {
+    timeoutId = setTimeout(() => {
       controller.abort("Request timeout");
     }, timeoutMs);
     try {
@@ -355,7 +372,7 @@ var apiCall = async (api, json) => {
       if (res.status === 401 || res.status === 403) {
         throw new AxAIServiceAuthenticationError(apiUrl.href, json, { metrics });
       }
-      if (res.status >= 400 && attempt < retryConfig.maxRetries && retryConfig.retryableStatusCodes.includes(res.status)) {
+      if (res.status >= 400 && shouldRetry(new Error(), res.status, attempt, retryConfig)) {
         const delay = calculateRetryDelay(attempt, retryConfig);
         attempt++;
         updateRetryMetrics(metrics);
@@ -369,7 +386,7 @@ var apiCall = async (api, json) => {
             "metrics.lastRetryTime": metrics.lastRetryTime
           });
         }
-        clearTimeout(timeoutId);
+        await new Promise((resolve) => setTimeout(resolve, delay));
         continue;
       }
       if (res.status >= 400) {
@@ -383,6 +400,17 @@ var apiCall = async (api, json) => {
       }
       if (!api.stream) {
         const resJson = await res.json();
+        if (api.validateResponse) {
+          const isValid = await api.validateResponse(resJson);
+          if (!isValid) {
+            throw new AxAIServiceResponseError(
+              "Invalid response data",
+              apiUrl.href,
+              json,
+              { validation: "response" }
+            );
+          }
+        }
         if (api.span?.isRecording()) {
           api.span.setAttributes({
             "response.time": Date.now() - metrics.startTime,
@@ -401,7 +429,7 @@ var apiCall = async (api, json) => {
       }
       let lastChunk;
       let chunkCount = 0;
-      const trackingStream = new TransformStream({
+      const trackingStream = new TransformStream3({
         transform(chunk, controller2) {
           lastChunk = chunk;
           chunkCount++;
@@ -409,7 +437,7 @@ var apiCall = async (api, json) => {
           metrics.lastChunkTime = Date.now();
           controller2.enqueue(chunk);
         },
-        flush(controller2) {
+        flush() {
           if (api.span?.isRecording()) {
             api.span.setAttributes({
               "stream.chunks": chunkCount,
@@ -417,10 +445,10 @@ var apiCall = async (api, json) => {
               "response.retries": metrics.retryCount
             });
           }
-          controller2.terminate();
         }
       });
-      const wrappedStream = new ReadableStream({
+      let closed = false;
+      return new ReadableStream({
         start(controller2) {
           const reader = res.body.pipeThrough(new textDecoderStream()).pipeThrough(new SSEParser()).pipeThrough(trackingStream).getReader();
           async function read() {
@@ -428,11 +456,14 @@ var apiCall = async (api, json) => {
               while (true) {
                 const { done, value } = await reader.read();
                 if (done) {
-                  controller2.close();
+                  if (!closed) {
+                    closed = true;
+                    controller2.close();
+                  }
                   break;
-                } else {
-                  controller2.enqueue(value);
                 }
+                if (closed) break;
+                controller2.enqueue(value);
               }
             } catch (e) {
               const error = e;
@@ -449,26 +480,42 @@ var apiCall = async (api, json) => {
                     { streamMetrics }
                   )
                 );
-              } else {
+              } else if (error instanceof TypeError && error.message.includes("cancelled")) {
                 controller2.error(
-                  new AxAIServiceResponseError(
-                    `Stream processing error: ${error.message}`,
+                  new AxAIServiceStreamTerminatedError(
                     apiUrl.href,
                     json,
-                    { streamMetrics }
+                    lastChunk,
+                    {
+                      streamMetrics,
+                      cancelReason: "Stream cancelled by client"
+                    }
                   )
                 );
+              } else {
+                controller2.error(
+                  new AxAIServiceNetworkError(error, apiUrl.href, json, {
+                    streamMetrics
+                  })
+                );
               }
+              throw error;
             } finally {
+              clearTimeout(timeoutId);
               reader.releaseLock();
+              if (api.span?.isRecording()) {
+                api.span.end();
+              }
             }
           }
           read();
+        },
+        // When the consumer cancels the stream, set our flag to stop processing further.
+        cancel() {
+          closed = true;
         }
       });
-      return wrappedStream;
     } catch (error) {
-      clearTimeout(timeoutId);
       if (error instanceof Error && error.name === "AbortError") {
         throw new AxAIServiceTimeoutError(apiUrl.href, timeoutMs, json, {
           metrics
@@ -481,7 +528,7 @@ var apiCall = async (api, json) => {
           "error.retries": metrics.retryCount
         });
       }
-      if (error instanceof AxAIServiceNetworkError && attempt < retryConfig.maxRetries) {
+      if (error instanceof AxAIServiceNetworkError && shouldRetry(error, void 0, attempt, retryConfig)) {
         const delay = calculateRetryDelay(attempt, retryConfig);
         attempt++;
         updateRetryMetrics(metrics);
@@ -495,12 +542,20 @@ var apiCall = async (api, json) => {
             "metrics.lastRetryTime": metrics.lastRetryTime
           });
         }
+        await new Promise((resolve) => setTimeout(resolve, delay));
         continue;
       }
       if (error instanceof AxAIServiceError) {
         error.context["metrics"] = metrics;
       }
       throw error;
+    } finally {
+      if (timeoutId !== void 0) {
+        clearTimeout(timeoutId);
+      }
+      if (api.span?.isRecording()) {
+        api.span.end();
+      }
     }
   }
 };
@@ -1057,8 +1112,6 @@ import { GoogleAuth } from "google-auth-library";
 var GoogleVertexAuth = class {
   auth;
   client;
-  currentToken;
-  tokenExpiry;
   constructor(config = {}) {
     this.auth = new GoogleAuth({
       scopes: ["https://www.googleapis.com/auth/cloud-platform"],
@@ -1072,36 +1125,12 @@ var GoogleVertexAuth = class {
     return this.client;
   }
   async getAccessToken() {
-    if (this.currentToken && this.tokenExpiry && Date.now() < this.tokenExpiry) {
-      return this.currentToken;
-    }
     const client = await this.getAuthenticatedClient();
-    const tokenResponse = await client.getAccessToken();
-    this.currentToken = tokenResponse.token;
-    const expiry = this.getExpiry(tokenResponse);
-    const tenMinutes = 10 * 60 * 1e3;
-    this.tokenExpiry = expiry - tenMinutes;
-    return this.currentToken;
-  }
-  /**
-   * Get the expiry date from the token response.
-   */
-  getExpiry(tokenResponse) {
-    const oneHour = 3600 * 1e3;
-    let expiry = Date.now() + oneHour;
-    let responseExpiry = tokenResponse.res?.data?.expiry_date;
-    if (responseExpiry) {
-      if (typeof responseExpiry === "number") {
-        expiry = responseExpiry;
-      } else if (responseExpiry instanceof Date) {
-        expiry = responseExpiry.getTime();
-      } else if (typeof responseExpiry === "string") {
-        expiry = new Date(responseExpiry).getTime();
-      } else {
-        console.warn("Unknown expiry type", responseExpiry);
-      }
+    const response = await client.getAccessToken();
+    if (!response.token) {
+      throw new Error("Failed to obtain access token");
     }
-    return expiry;
+    return response.token;
   }
 };
 
@@ -1578,8 +1607,9 @@ function mapFinishReason(stopReason) {
 
 // ai/openai/types.ts
 var AxAIOpenAIModel = /* @__PURE__ */ ((AxAIOpenAIModel2) => {
-  AxAIOpenAIModel2["O1Preview"] = "o1-preview";
+  AxAIOpenAIModel2["O1"] = "o1";
   AxAIOpenAIModel2["O1Mini"] = "o1-mini";
+  AxAIOpenAIModel2["O3Mini"] = "o3-mini";
   AxAIOpenAIModel2["GPT4"] = "gpt-4";
   AxAIOpenAIModel2["GPT4O"] = "gpt-4o";
   AxAIOpenAIModel2["GPT4OMini"] = "gpt-4o-mini";
@@ -1602,7 +1632,7 @@ var AxAIOpenAIEmbedModel = /* @__PURE__ */ ((AxAIOpenAIEmbedModel2) => {
 // ai/openai/info.ts
 var axModelInfoOpenAI = [
   {
-    name: "o1-preview" /* O1Preview */,
+    name: "o1" /* O1 */,
     currency: "usd",
     promptTokenCostPer1M: 15,
     completionTokenCostPer1M: 60
@@ -1610,8 +1640,14 @@ var axModelInfoOpenAI = [
   {
     name: "o1-mini" /* O1Mini */,
     currency: "usd",
-    promptTokenCostPer1M: 3,
-    completionTokenCostPer1M: 12
+    promptTokenCostPer1M: 1.1,
+    completionTokenCostPer1M: 14.4
+  },
+  {
+    name: "o3-mini" /* O3Mini */,
+    currency: "usd",
+    promptTokenCostPer1M: 1.1,
+    completionTokenCostPer1M: 4.4
   },
   {
     name: "gpt-4" /* GPT4 */,
@@ -1711,16 +1747,12 @@ var AxAIOpenAIImpl = class {
         parameters: v.parameters
       }
     }));
-    if (tools && isO1Model(model)) {
-      throw new Error("Functions are not supported for O1 models");
-    }
     const toolsChoice = !req.functionCall && req.functions && req.functions.length > 0 ? "auto" : req.functionCall;
     const messages = createMessages2(req);
     const frequencyPenalty = req.modelConfig?.frequencyPenalty ?? this.config.frequencyPenalty;
     const stream = req.modelConfig?.stream ?? this.config.stream;
-    if (stream && isO1Model(model)) {
-      throw new Error("Streaming is not supported for O1 models");
-    }
+    const reasoningEffort = isReasoningModel(model) ? this.config.reasoningEffort : void 0;
+    const store = this.config.store;
     const reqValue = {
       model,
       messages,
@@ -1735,7 +1767,9 @@ var AxAIOpenAIImpl = class {
       presence_penalty: req.modelConfig?.presencePenalty ?? this.config.presencePenalty,
       logit_bias: this.config.logitBias,
       ...frequencyPenalty ? { frequency_penalty: frequencyPenalty } : {},
-      ...stream && this.streamingUsage ? { stream: true, stream_options: { include_usage: true } } : {}
+      ...stream && this.streamingUsage ? { stream: true, stream_options: { include_usage: true } } : {},
+      ...reasoningEffort ? { reasoning_effort: reasoningEffort } : {},
+      ...store ? { store } : {}
     };
     return [apiConfig, reqValue];
   }
@@ -1861,12 +1895,6 @@ var mapFinishReason2 = (finishReason) => {
 };
 function createMessages2(req) {
   return req.chatPrompt.map((msg) => {
-    if (msg.role === "system" && isO1Model(req.model)) {
-      msg = {
-        role: "user",
-        content: msg.content
-      };
-    }
     switch (msg.role) {
       case "system":
         return { role: "system", content: msg.content };
@@ -1956,14 +1984,14 @@ var AxAIOpenAI = class extends AxBaseAI {
         embedModel: _config.embedModel
       },
       options,
-      supportFor: (model) => {
-        return isO1Model(model) ? { functions: false, streaming: false } : { functions: true, streaming: true };
+      supportFor: () => {
+        return { functions: true, streaming: true };
       },
       modelMap
     });
   }
 };
-var isO1Model = (model) => ["o1-mini" /* O1Mini */, "o1-preview" /* O1Preview */].includes(
+var isReasoningModel = (model) => ["o1-mini" /* O1Mini */, "o1" /* O1 */, "o3-mini" /* O3Mini */].includes(
   model
 );
 
@@ -5208,8 +5236,10 @@ function* streamValues(sig, values, xstate, content, final = false) {
       const v = content.substring(s);
       const v1 = v.replace(/[\s\n\t]+$/, "");
       const v2 = pos === 0 ? v1.trimStart() : v1;
-      yield { [fieldName]: v2 };
-      xstate.streamedIndex[fieldName] = pos + v1.length;
+      if (v2.length > 0) {
+        yield { [fieldName]: v2 };
+        xstate.streamedIndex[fieldName] = pos + v1.length;
+      }
       return;
     }
   }
@@ -5218,7 +5248,7 @@ function* streamValues(sig, values, xstate, content, final = false) {
     if (Array.isArray(value)) {
       const s = xstate.streamedIndex[key] ?? 0;
       const v = value.slice(s);
-      if (v) {
+      if (v && v.length > 0) {
         yield { [key]: v };
         xstate.streamedIndex[key] = s + 1;
       }
@@ -5314,7 +5344,7 @@ var extractBlock = (input) => {
 // dsp/jsonschema.ts
 var validateJSONSchema = (schema) => {
   const errors = [];
-  const validateSchemaObject = (schema2, path2 = "") => {
+  const validateSchemaObject = (schema2, path = "") => {
     const validTypes = [
       "array",
       "integer",
@@ -5325,31 +5355,31 @@ var validateJSONSchema = (schema) => {
       "object"
     ];
     if (!validTypes.includes(schema2.type)) {
-      errors.push(`Invalid type '${schema2.type}' at ${path2 || "root"}`);
+      errors.push(`Invalid type '${schema2.type}' at ${path || "root"}`);
       return;
     }
     if (schema2.type === "object" && schema2.properties) {
       if (typeof schema2.properties !== "object" || Array.isArray(schema2.properties)) {
-        errors.push(`Invalid properties object at ${path2 || "root"}`);
+        errors.push(`Invalid properties object at ${path || "root"}`);
       } else {
         for (const key in schema2.properties) {
           const value = schema2.properties[key];
           if (typeof value !== "object") {
-            errors.push(`Invalid schema object at ${path2}${key}`);
+            errors.push(`Invalid schema object at ${path}${key}`);
             continue;
           }
-          validateSchemaObject(value, `${path2}${key}.`);
+          validateSchemaObject(value, `${path}${key}.`);
         }
       }
       if (schema2.required && !Array.isArray(schema2.required)) {
-        errors.push(`'required' should be an array at ${path2 || "root"}`);
+        errors.push(`'required' should be an array at ${path || "root"}`);
       }
     }
     if (schema2.type === "array" && schema2.items) {
       if (typeof schema2.items !== "object") {
-        errors.push(`Invalid items schema at ${path2 || "root"}`);
+        errors.push(`Invalid items schema at ${path || "root"}`);
       } else {
-        validateSchemaObject(schema2.items, `${path2}items.`);
+        validateSchemaObject(schema2.items, `${path}items.`);
       }
     }
   };
@@ -5733,6 +5763,7 @@ var AxGen = class extends AxProgramWithSignature {
             const e1 = e;
             errorFields = e1.getFixingInstructions();
             err = e;
+          } else if (e instanceof AxAIServiceStreamTerminatedError) {
           } else {
             throw e;
           }
@@ -5996,6 +6027,7 @@ var AxBalancer = class _AxBalancer {
   services;
   currentServiceIndex = 0;
   currentService;
+  debug;
   constructor(services, options) {
     if (services.length === 0) {
       throw new Error("No AI services provided.");
@@ -6008,6 +6040,7 @@ var AxBalancer = class _AxBalancer {
       throw new Error("Error initializing the AI services.");
     }
     this.currentService = cs;
+    this.debug = options?.debug ?? true;
   }
   /**
    * Service comparator that respects the input order of services.
@@ -6082,16 +6115,20 @@ var AxBalancer = class _AxBalancer {
           default:
             throw e;
         }
-        console.warn(
-          `AxBalancer: Service ${this.currentService.getName()} failed`,
-          e
-        );
+        if (this.debug) {
+          console.warn(
+            `AxBalancer: Service ${this.currentService.getName()} failed`,
+            e
+          );
+        }
         if (!this.getNextService()) {
           throw e;
         }
-        console.warn(
-          `AxBalancer: Switching to service ${this.currentService.getName()}`
-        );
+        if (this.debug) {
+          console.warn(
+            `AxBalancer: Switching to service ${this.currentService.getName()}`
+          );
+        }
       }
     }
   }
@@ -6101,11 +6138,15 @@ var AxBalancer = class _AxBalancer {
       try {
         return await this.currentService.embed(req, options);
       } catch (e) {
-        console.warn(`Service ${this.currentService.getName()} failed`);
+        if (this.debug) {
+          console.warn(`Service ${this.currentService.getName()} failed`);
+        }
         if (!this.getNextService()) {
           throw e;
         }
-        console.warn(`Switching to service ${this.currentService.getName()}`);
+        if (this.debug) {
+          console.warn(`Switching to service ${this.currentService.getName()}`);
+        }
       }
     }
   }