@ax-llm/ax 10.0.43 → 10.0.45

package/index.js

@@ -76,22 +76,29 @@ var AxSpanKindValues = /* @__PURE__ */ ((AxSpanKindValues2) => {
 // util/apicall.ts
 import path from "path";
 import {
+  ReadableStream,
   TextDecoderStream as TextDecoderStreamNative
 } from "stream/web";
 import "@opentelemetry/api";
 
 // util/sse.ts
-import { TransformStream } from "stream/web";
-var SSEParser = class extends TransformStream {
-  constructor(dataParser = JSON.parse) {
+import { TransformStream as TransformStream2 } from "stream/web";
+var SSEParser = class extends TransformStream2 {
+  buffer = "";
+  currentEvent = { rawData: "" };
+  dataParser;
+  onError;
+  constructor(options = {}) {
     super({
       transform: (chunk, controller) => this.handleChunk(chunk, controller),
       flush: (controller) => this.handleFlush(controller)
     });
-    this.dataParser = dataParser;
+    this.dataParser = options.dataParser || JSON.parse;
+    this.onError = options.onError || ((error, rawData) => {
+      console.warn("Failed to parse event data:", error);
+      console.log("Raw data that failed to parse:", rawData);
+    });
   }
-  buffer = "";
-  currentEvent = { rawData: "" };
   handleChunk(chunk, controller) {
     this.buffer += chunk;
     this.processBuffer(controller);
@@ -99,24 +106,28 @@ var SSEParser = class extends TransformStream {
   handleFlush(controller) {
     this.processBuffer(controller);
     if (this.currentEvent.rawData) {
-      this.emitEvent(controller);
+      this.processEvent(controller);
     }
   }
   processBuffer(controller) {
-    const lines = this.buffer.split(/\r\n|\r|\n/);
+    const normalizedBuffer = this.buffer.replace(/\r\n|\r/g, "\n");
+    const lines = normalizedBuffer.split("\n");
     this.buffer = lines.pop() || "";
     for (const line of lines) {
-      if (line.trim() === "") {
-        this.emitEvent(controller);
+      if (line === "") {
+        this.processEvent(controller);
       } else {
         this.parseLine(line);
       }
     }
   }
   parseLine(line) {
+    if (line.startsWith(":")) {
+      return;
+    }
     const colonIndex = line.indexOf(":");
     if (colonIndex === -1) {
-      this.currentEvent.rawData += this.currentEvent.rawData ? "\n" + line.trim() : line.trim();
+      this.currentEvent.rawData += (this.currentEvent.rawData && !this.currentEvent.rawData.endsWith("\n") ? "\n" : "") + line.trim();
       return;
     }
     const field = line.slice(0, colonIndex).trim();
@@ -126,7 +137,7 @@ var SSEParser = class extends TransformStream {
         this.currentEvent.event = value;
         break;
       case "data":
-        this.currentEvent.rawData += this.currentEvent.rawData ? "\n" + value : value;
+        this.currentEvent.rawData += (this.currentEvent.rawData && !this.currentEvent.rawData.endsWith("\n") ? "\n" : "") + value;
         break;
       case "id":
         this.currentEvent.id = value;
@@ -140,21 +151,20 @@ var SSEParser = class extends TransformStream {
       }
     }
   }
-  emitEvent(controller) {
+  processEvent(controller) {
     if (this.currentEvent.rawData) {
-      if (this.currentEvent.rawData.trim() === "[DONE]" || this.currentEvent.rawData.trim().startsWith("[")) {
+      if (!this.currentEvent.event) {
+        this.currentEvent.event = "message";
+      }
+      if (this.currentEvent.rawData.trim() === "[DONE]") {
+        this.currentEvent = { rawData: "" };
         return;
-      } else {
-        try {
-          const parsedData = this.dataParser(this.currentEvent.rawData);
-          controller.enqueue(parsedData);
-        } catch (e) {
-          console.warn("Failed to parse event data:", e);
-          console.log(
-            "Raw data that failed to parse:",
-            this.currentEvent.rawData
-          );
-        }
+      }
+      try {
+        const parsedData = this.dataParser(this.currentEvent.rawData);
+        controller.enqueue(parsedData);
+      } catch (e) {
+        this.onError(e, this.currentEvent.rawData);
       }
       this.currentEvent = { rawData: "" };
     }
@@ -163,7 +173,7 @@ var SSEParser = class extends TransformStream {
 
 // util/stream.ts
 import {
-  TransformStream as TransformStream2
+  TransformStream as TransformStream3
 } from "stream/web";
 var TextDecodeTransformer = class {
   decoder;
@@ -186,59 +196,312 @@ var TextDecodeTransformer = class {
     }
   }
 };
-var TextDecoderStreamPolyfill = class extends TransformStream2 {
+var TextDecoderStreamPolyfill = class extends TransformStream3 {
   constructor() {
     super(new TextDecodeTransformer());
   }
 };
 
 // util/apicall.ts
+var defaultRetryConfig = {
+  maxRetries: 3,
+  initialDelayMs: 1e3,
+  maxDelayMs: 6e4,
+  // Increased to 60 seconds
+  backoffFactor: 2,
+  retryableStatusCodes: [500, 408, 429, 502, 503, 504]
+};
+var defaultTimeoutMs = 3e4;
 var textDecoderStream = TextDecoderStreamNative ?? TextDecoderStreamPolyfill;
+var AxAIServiceError = class extends Error {
+  constructor(message, url, requestBody, context = {}) {
+    super(message);
+    this.url = url;
+    this.requestBody = requestBody;
+    this.context = context;
+    this.name = "AxAIServiceError";
+    this.timestamp = (/* @__PURE__ */ new Date()).toISOString();
+    this.errorId = crypto.randomUUID();
+  }
+  timestamp;
+  errorId;
+  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 ? `
+Context: ${JSON.stringify(this.context, null, 2)}` : ""}`;
+  }
+  toJSON() {
+    return {
+      name: this.name,
+      errorId: this.errorId,
+      message: this.message,
+      timestamp: this.timestamp,
+      url: this.url,
+      requestBody: this.requestBody,
+      context: this.context,
+      stack: this.stack
+    };
+  }
+};
+var AxAIServiceStatusError = class extends AxAIServiceError {
+  constructor(status, statusText, url, requestBody, context) {
+    super(`HTTP ${status} - ${statusText}`, url, requestBody, {
+      httpStatus: status,
+      httpStatusText: statusText,
+      ...context
+    });
+    this.status = status;
+    this.statusText = statusText;
+    this.name = "AxAIServiceStatusError";
+  }
+};
+var AxAIServiceNetworkError = class extends AxAIServiceError {
+  constructor(originalError, url, requestBody, context) {
+    super(`Network Error: ${originalError.message}`, url, requestBody, {
+      originalErrorName: originalError.name,
+      originalErrorStack: originalError.stack,
+      ...context
+    });
+    this.originalError = originalError;
+    this.name = "AxAIServiceNetworkError";
+    this.stack = originalError.stack;
+  }
+};
+var AxAIServiceResponseError = class extends AxAIServiceError {
+  constructor(message, url, requestBody, context) {
+    super(message, url, requestBody, context);
+    this.name = "AxAIServiceResponseError";
+  }
+};
+var AxAIServiceStreamTerminatedError = class extends AxAIServiceError {
+  constructor(url, requestBody, lastChunk, context) {
+    super("Stream terminated unexpectedly by remote host", url, requestBody, {
+      lastChunk,
+      ...context
+    });
+    this.lastChunk = lastChunk;
+    this.name = "AxAIServiceStreamTerminatedError";
+  }
+};
+var AxAIServiceTimeoutError = class extends AxAIServiceError {
+  constructor(url, timeoutMs, requestBody, context) {
+    super(`Request timeout after ${timeoutMs}ms`, url, requestBody, {
+      timeoutMs,
+      ...context
+    });
+    this.name = "AxAIServiceTimeoutError";
+  }
+};
+var AxAIServiceAuthenticationError = class extends AxAIServiceError {
+  constructor(url, requestBody, context) {
+    super("Authentication failed", url, requestBody, context);
+    this.name = "AxAIServiceAuthenticationError";
+  }
+};
+function calculateRetryDelay(attempt, config) {
+  const delay = Math.min(
+    config.maxDelayMs,
+    config.initialDelayMs * Math.pow(config.backoffFactor, attempt)
+  );
+  return delay * (0.75 + Math.random() * 0.5);
+}
+function createRequestMetrics() {
+  return {
+    startTime: Date.now(),
+    retryCount: 0
+  };
+}
+function updateRetryMetrics(metrics) {
+  metrics.retryCount++;
+  metrics.lastRetryTime = Date.now();
+}
 var apiCall = async (api, json) => {
+  const retryConfig = { ...defaultRetryConfig, ...api.retry };
+  const timeoutMs = api.timeout ?? defaultTimeoutMs;
+  const metrics = createRequestMetrics();
   const baseUrl = new URL(process.env["PROXY"] ?? api.url);
   const apiPath = path.join(baseUrl.pathname, api.name ?? "/", baseUrl.search);
   const apiUrl = new URL(apiPath, baseUrl);
+  const requestId = crypto.randomUUID();
   if (api.span?.isRecording()) {
     api.span.setAttributes({
       "http.request.method": api.put ? "PUT" : "POST",
-      "url.full": apiUrl.href
+      "url.full": apiUrl.href,
+      "request.id": requestId,
+      "request.startTime": metrics.startTime
     });
   }
-  let res;
-  try {
-    res = await (api.fetch ?? fetch)(apiUrl, {
-      method: api.put ? "PUT" : "POST",
-      headers: {
-        "Content-Type": "application/json",
-        ...api.headers
-      },
-      body: JSON.stringify(json)
-    });
-    if (res.status >= 400) {
-      const reqBody = JSON.stringify(json, null, 2);
-      throw new Error(
-        `API Request Error: ${res.status}, ${res.statusText}
-:Request Body: ${reqBody}`
-      );
-    }
-    if (!api.stream) {
-      const resJson = await res.json();
-      return resJson;
-    }
-    if (!res.body) {
-      throw new Error("Response body is null");
-    }
-    const st = res.body.pipeThrough(new textDecoderStream()).pipeThrough(new SSEParser());
-    return st;
-  } catch (e) {
-    if (api.span?.isRecording()) {
-      api.span.recordException(e);
+  let attempt = 0;
+  while (true) {
+    const controller = new AbortController();
+    let timeoutId = setTimeout(() => {
+      controller.abort("Request timeout");
+    }, timeoutMs);
+    try {
+      const res = await (api.fetch ?? fetch)(apiUrl, {
+        method: api.put ? "PUT" : "POST",
+        headers: {
+          "Content-Type": "application/json",
+          "X-Request-ID": requestId,
+          "X-Retry-Count": attempt.toString(),
+          ...api.headers
+        },
+        body: JSON.stringify(json),
+        signal: controller.signal
+      });
+      clearTimeout(timeoutId);
+      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)) {
+        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
+          });
+        }
+        clearTimeout(timeoutId);
+        continue;
+      }
+      if (res.status >= 400) {
+        throw new AxAIServiceStatusError(
+          res.status,
+          res.statusText,
+          apiUrl.href,
+          json,
+          { metrics }
+        );
+      }
+      if (!api.stream) {
+        const resJson = await res.json();
+        if (api.span?.isRecording()) {
+          api.span.setAttributes({
+            "response.time": Date.now() - metrics.startTime,
+            "response.retries": metrics.retryCount
+          });
+        }
+        return resJson;
+      }
+      if (!res.body) {
+        throw new AxAIServiceResponseError(
+          "Response body is null",
+          apiUrl.href,
+          json,
+          { metrics }
+        );
+      }
+      let lastChunk;
+      let chunkCount = 0;
+      const trackingStream = new TransformStream({
+        transform(chunk, controller2) {
+          lastChunk = chunk;
+          chunkCount++;
+          metrics.streamChunks = chunkCount;
+          metrics.lastChunkTime = Date.now();
+          controller2.enqueue(chunk);
+        },
+        flush(controller2) {
+          if (api.span?.isRecording()) {
+            api.span.setAttributes({
+              "stream.chunks": chunkCount,
+              "stream.duration": Date.now() - metrics.startTime,
+              "response.retries": metrics.retryCount
+            });
+          }
+          controller2.terminate();
+        }
+      });
+      const wrappedStream = new ReadableStream({
+        start(controller2) {
+          const reader = res.body.pipeThrough(new textDecoderStream()).pipeThrough(new SSEParser()).pipeThrough(trackingStream).getReader();
+          async function read() {
+            try {
+              while (true) {
+                const { done, value } = await reader.read();
+                if (done) {
+                  controller2.close();
+                  break;
+                } else {
+                  controller2.enqueue(value);
+                }
+              }
+            } catch (e) {
+              const error = e;
+              const streamMetrics = {
+                ...metrics,
+                streamDuration: Date.now() - metrics.startTime
+              };
+              if (error.name === "AbortError" || error.message?.includes("aborted")) {
+                controller2.error(
+                  new AxAIServiceStreamTerminatedError(
+                    apiUrl.href,
+                    json,
+                    lastChunk,
+                    { streamMetrics }
+                  )
+                );
+              } else {
+                controller2.error(
+                  new AxAIServiceResponseError(
+                    `Stream processing error: ${error.message}`,
+                    apiUrl.href,
+                    json,
+                    { streamMetrics }
+                  )
+                );
+              }
+            } finally {
+              reader.releaseLock();
+            }
+          }
+          read();
+        }
+      });
+      return wrappedStream;
+    } catch (error) {
+      clearTimeout(timeoutId);
+      if (error instanceof Error && error.name === "AbortError") {
+        throw new AxAIServiceTimeoutError(apiUrl.href, timeoutMs, json, {
+          metrics
+        });
+      }
+      if (api.span?.isRecording()) {
+        api.span.recordException(error);
+        api.span.setAttributes({
+          "error.time": Date.now() - metrics.startTime,
+          "error.retries": metrics.retryCount
+        });
+      }
+      if (error instanceof AxAIServiceNetworkError && attempt < retryConfig.maxRetries) {
+        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
+          });
+        }
+        continue;
+      }
+      if (error instanceof AxAIServiceError) {
+        error.context["metrics"] = metrics;
+      }
+      throw error;
     }
-    const reqBody = JSON.stringify(json, null, 2);
-    throw new Error(
-      `API Response Error: ${apiUrl.href}, ${e}
-Request Body: ${reqBody}`
-    );
   }
 };
 
@@ -275,7 +538,7 @@ var ColorLog = class {
 
 // util/transform.ts
 import {
-  TransformStream as TransformStream3
+  TransformStream as TransformStream4
 } from "stream/web";
 var TypeTransformer = class {
   buffer;
@@ -298,7 +561,7 @@ var TypeTransformer = class {
     controller.terminate();
   }
 };
-var RespTransformStream = class extends TransformStream3 {
+var RespTransformStream = class extends TransformStream4 {
   constructor(transformFn, doneCallback) {
     super(new TypeTransformer(transformFn, doneCallback));
   }
@@ -814,10 +1077,10 @@ var GoogleVertexAuth = class {
     }
     const client = await this.getAuthenticatedClient();
     const tokenResponse = await client.getAccessToken();
-    this.currentToken = tokenResponse.token ?? void 0;
+    this.currentToken = tokenResponse.token;
     const expiry = this.getExpiry(tokenResponse);
-    const fiveMinutes = 5 * 60 * 1e3;
-    this.tokenExpiry = expiry - fiveMinutes;
+    const tenMinutes = 10 * 60 * 1e3;
+    this.tokenExpiry = expiry - tenMinutes;
     return this.currentToken;
   }
   /**
@@ -837,8 +1100,6 @@ var GoogleVertexAuth = class {
       } else {
         console.warn("Unknown expiry type", responseExpiry);
       }
-    } else {
-      console.warn("No expiry date found in response", tokenResponse.res?.data);
     }
     return expiry;
   }
@@ -3261,7 +3522,7 @@ var AxAI = class {
 };
 
 // dsp/generate.ts
-import { ReadableStream } from "stream/web";
+import { ReadableStream as ReadableStream2 } from "stream/web";
 import { SpanKind as SpanKind2 } from "@opentelemetry/api";
 
 // ai/util.ts
@@ -5276,7 +5537,7 @@ var AxGen = class extends AxProgramWithSignature {
       mem,
       options
     });
-    if (res instanceof ReadableStream) {
+    if (res instanceof ReadableStream2) {
       yield* this.processStreamingResponse({
         ai,
         model,
@@ -5802,11 +6063,35 @@ var AxBalancer = class _AxBalancer {
       try {
         return await this.currentService.chat(req, options);
       } catch (e) {
-        console.warn(`Service ${this.currentService.getName()} failed`);
+        if (!(e instanceof AxAIServiceError)) {
+          throw e;
+        }
+        switch (e.constructor) {
+          case AxAIServiceAuthenticationError:
+            throw e;
+          case AxAIServiceStatusError:
+            break;
+          case AxAIServiceNetworkError:
+            break;
+          case AxAIServiceResponseError:
+            break;
+          case AxAIServiceStreamTerminatedError:
+            break;
+          case AxAIServiceTimeoutError:
+            break;
+          default:
+            throw e;
+        }
+        console.warn(
+          `AxBalancer: Service ${this.currentService.getName()} failed`,
+          e
+        );
         if (!this.getNextService()) {
           throw e;
         }
-        console.warn(`Switching to service ${this.currentService.getName()}`);
+        console.warn(
+          `AxBalancer: Switching to service ${this.currentService.getName()}`
+        );
       }
     }
   }
@@ -7487,6 +7772,13 @@ export {
   AxAIOpenAIModel,
   AxAIReka,
   AxAIRekaModel,
+  AxAIServiceAuthenticationError,
+  AxAIServiceError,
+  AxAIServiceNetworkError,
+  AxAIServiceResponseError,
+  AxAIServiceStatusError,
+  AxAIServiceStreamTerminatedError,
+  AxAIServiceTimeoutError,
   AxAITogether,
   AxAgent,
   AxApacheTika,