@ax-llm/ax 11.0.57 → 11.0.58

package/index.js

@@ -669,6 +669,46 @@ var ColorLog = class {
 
 // ai/debug.ts
 var colorLog = new ColorLog();
+var defaultOutput = (message) => {
+  process.stdout.write(message);
+};
+var createDefaultLogger = (output = defaultOutput) => {
+  return (message, options) => {
+    const tags = options?.tags ?? [];
+    let formattedMessage = message;
+    if (tags.includes("error")) {
+      formattedMessage = colorLog.red(formattedMessage);
+    } else if (tags.includes("success") || tags.includes("responseContent")) {
+      formattedMessage = colorLog.greenBright(formattedMessage);
+    } else if (tags.includes("functionName")) {
+      formattedMessage = colorLog.whiteBright(formattedMessage);
+    } else if (tags.includes("functionArg") || tags.includes("systemContent") || tags.includes("assistantContent")) {
+      formattedMessage = colorLog.blueBright(formattedMessage);
+    } else if (tags.includes("warning") || tags.includes("discovery")) {
+      formattedMessage = colorLog.yellow(formattedMessage);
+    }
+    if (tags.includes("responseStart") || tags.includes("systemStart") || tags.includes("userStart")) {
+      formattedMessage = `
+${formattedMessage}`;
+    } else if (tags.includes("responseEnd") || tags.includes("systemEnd") || tags.includes("userEnd")) {
+      formattedMessage = `${formattedMessage}
+`;
+    } else if (tags.includes("assistantStart")) {
+      formattedMessage = `
+${formattedMessage}
+`;
+    } else if (tags.includes("error")) {
+      formattedMessage = `
+${formattedMessage}
+`;
+    } else if (tags.includes("functionEnd")) {
+      formattedMessage = `${formattedMessage}
+`;
+    }
+    output(formattedMessage);
+  };
+};
+var defaultLogger = createDefaultLogger();
 var formatChatMessage = (msg, hideContent, hideSystemPrompt) => {
   switch (msg.role) {
     case "system":
@@ -676,30 +716,30 @@ var formatChatMessage = (msg, hideContent, hideSystemPrompt) => {
         return "";
       }
       return `
-${colorLog.blueBright("System:")}
-${colorLog.whiteBright(msg.content)}`;
+System:
+${msg.content}`;
     case "function":
       return `
-${colorLog.blueBright("Function Result:")}
-${colorLog.whiteBright(msg.result)}`;
+Function Result:
+${msg.result}`;
     case "user": {
       if (typeof msg.content === "string") {
         return `
-${colorLog.blueBright("User:")}
-${colorLog.whiteBright(msg.content)}`;
+User:
+${msg.content}`;
       }
       const items = msg.content.map((v) => {
         switch (v.type) {
           case "text":
-            return `${colorLog.whiteBright(v.text)}`;
+            return v.text;
           case "image":
-            return `(Image, ${v.mimeType}) ${colorLog.whiteBright(v.image.substring(0, 10))}`;
+            return `(Image, ${v.mimeType}) ${v.image.substring(0, 10)}`;
           default:
             throw new Error("Invalid content type");
         }
       });
       return `
-${colorLog.blueBright("User:")}
+User:
 ${items.join("\n")}`;
     }
     case "assistant": {
@@ -709,62 +749,64 @@ ${items.join("\n")}`;
           return `${fn.name}(${args})`;
         });
         return `
-${colorLog.blueBright("\nFunctions:")}
-${colorLog.whiteBright(fns.join("\n"))}`;
+Functions:
+${fns.join("\n")}`;
       }
       return `
-${colorLog.blueBright("\nAssistant:")}
-${hideContent ? "" : colorLog.whiteBright(msg.content ?? "<empty>")}`;
+Assistant:
+${hideContent ? "" : msg.content ?? "<empty>"}`;
     }
     default:
       throw new Error("Invalid role");
   }
 };
-var logChatRequestMessage = (msg, hideSystemPrompt) => {
-  process.stdout.write(`${formatChatMessage(msg, hideSystemPrompt)}
-`);
-  process.stdout.write(colorLog.blueBright("\nAssistant:\n"));
+var logChatRequestMessage = (msg, hideSystemPrompt, logger = defaultLogger) => {
+  const formattedMessage = formatChatMessage(msg, false, hideSystemPrompt);
+  if (formattedMessage) {
+    const tags = msg.role === "system" ? ["systemStart", "systemContent"] : msg.role === "function" ? ["functionName"] : msg.role === "user" ? ["userStart", "userContent"] : [];
+    logger(formattedMessage, { tags });
+  }
+  logger("Assistant:", { tags: ["assistantStart"] });
 };
-var logChatRequest = (chatPrompt, hideSystemPrompt) => {
-  const items = chatPrompt?.map(
-    (msg) => formatChatMessage(msg, hideSystemPrompt)
-  );
-  if (items) {
-    process.stdout.write(items.join("\n"));
-    process.stdout.write(colorLog.blueBright("\nAssistant:\n"));
+var logChatRequest = (chatPrompt, hideSystemPrompt, logger = defaultLogger) => {
+  for (const msg of chatPrompt ?? []) {
+    const formattedMessage = formatChatMessage(msg, false, hideSystemPrompt);
+    if (formattedMessage) {
+      const tags = msg.role === "system" ? ["systemContent"] : msg.role === "function" ? ["functionName"] : msg.role === "user" ? ["userContent"] : [];
+      logger(formattedMessage, { tags });
+    }
   }
+  logger("Assistant:", { tags: ["assistantStart"] });
 };
-var logResponseResult = (r) => {
+var logResponseResult = (r, logger = defaultLogger) => {
   if (r.content) {
-    process.stdout.write(colorLog.greenBright(r.content));
+    logger(r.content, { tags: ["responseContent"] });
   }
-  if (r.functionCalls) {
+  if (r.functionCalls && r.functionCalls.length > 0) {
     for (const [i, f] of r.functionCalls.entries()) {
       if (f.function.name) {
-        if (i > 0) {
-          process.stdout.write("\n");
-        }
-        process.stdout.write(
-          `Function ${i + 1} -> ${colorLog.greenBright(f.function.name)}`
-        );
+        logger(`[${i + 1}] ${f.function.name}`, {
+          tags: ["functionName"]
+        });
       }
       if (f.function.params) {
         const params = typeof f.function.params === "string" ? f.function.params : JSON.stringify(f.function.params, null, 2);
-        process.stdout.write(`${colorLog.greenBright(params)}`);
+        logger(params, { tags: ["functionArg"] });
       }
     }
+    logger("", { tags: ["functionEnd"] });
   }
 };
-var logResponse = (resp) => {
+var logResponse = (resp, logger = defaultLogger) => {
   if (!resp.results) {
     return;
   }
   for (const r of resp.results) {
-    logResponseResult(r);
+    logResponseResult(r, logger);
   }
 };
-var logResponseDelta = (delta) => {
-  process.stdout.write(colorLog.greenBright(delta));
+var logResponseDelta = (delta, logger = defaultLogger) => {
+  logger(delta, { tags: ["responseContent"] });
 };
 
 // ai/base.ts
@@ -778,6 +820,9 @@ var axBaseAIDefaultCreativeConfig = () => structuredClone({
   topP: 0.7,
   frequencyPenalty: 0.2
 });
+var defaultLogger2 = (message, _options) => {
+  process.stdout.write(message);
+};
 var AxBaseAI = class {
   constructor(aiImpl, {
     name,
@@ -817,6 +862,7 @@ var AxBaseAI = class {
   excludeContentFromTrace;
   models;
   abortSignal;
+  logger = defaultLogger2;
   modelInfo;
   modelUsage;
   embedModelUsage;
@@ -878,6 +924,7 @@ var AxBaseAI = class {
     this.tracer = options.tracer;
     this.excludeContentFromTrace = options.excludeContentFromTrace;
     this.abortSignal = options.abortSignal;
+    this.logger = options.logger ?? defaultLogger2;
   }
   getOptions() {
     return {
@@ -887,9 +934,13 @@ var AxBaseAI = class {
       tracer: this.tracer,
       timeout: this.timeout,
       excludeContentFromTrace: this.excludeContentFromTrace,
-      abortSignal: this.abortSignal
+      abortSignal: this.abortSignal,
+      logger: this.logger
     };
   }
+  getLogger() {
+    return this.logger;
+  }
   getModelList() {
     const models = [];
     for (const model of this.models ?? []) {
@@ -1077,7 +1128,11 @@ var AxBaseAI = class {
       return res2;
     };
     if (debug) {
-      logChatRequest(req.chatPrompt, options?.debugHideSystemPrompt);
+      logChatRequest(
+        req.chatPrompt,
+        options?.debugHideSystemPrompt,
+        options?.logger ?? this.logger
+      );
     }
     const rt = options?.rateLimiter ?? this.rt;
     const rv = rt ? await rt(fn, { modelUsage: this.modelUsage }) : await fn();
@@ -1108,13 +1163,14 @@ var AxBaseAI = class {
           setChatResponseEvents(res2, span, this.excludeContentFromTrace);
         }
         if (debug) {
-          logResponse(res2);
+          logResponse(res2, options?.logger ?? this.logger);
         }
         return res2;
       };
       const doneCb = async (_values) => {
         if (debug) {
-          process.stdout.write("\n");
+          const logger = options?.logger ?? this.logger;
+          logger("", { tags: ["responseEnd"] });
         }
         if (span?.isRecording()) {
           span.end();
@@ -1158,7 +1214,10 @@ var AxBaseAI = class {
       span.end();
     }
     if (debug) {
-      logResponse(res);
+      logResponse(res, options?.logger ?? this.logger);
+    }
+    if (debug) {
+      this.logger("", { tags: ["responseEnd"] });
     }
     return res;
   }
@@ -5143,6 +5202,9 @@ var AxAI = class {
   getOptions() {
     return this.ai.getOptions();
   }
+  getLogger() {
+    return this.ai.getLogger();
+  }
 };
 
 // ai/x-grok/types.ts
@@ -5366,7 +5428,7 @@ var MemoryImpl = class {
     if (this.options?.debug) {
       if (delta && typeof delta === "string") {
         debugResponseDelta(delta);
-      } else if (lastItem) {
+      } else if (!delta && (content || functionCalls)) {
         debugResponse({ content, name, functionCalls });
       }
     }
@@ -6308,7 +6370,6 @@ function formatDescription(str) {
 }
 
 // dsp/validate.ts
-var colorLog4 = new ColorLog();
 var ValidationError = class extends Error {
   fields;
   constructor({
@@ -6349,10 +6410,11 @@ function handleValidationError(mem, errorFields, ai, promptTemplate, sessionId)
   mem.addTag("error");
   if (ai.getOptions().debug) {
     const errors = errorFields.map((field) => `- ${field.title}: ${field.description}`).join("\n");
-    process.stdout.write(colorLog4.red(`
-\u274C Error Correction:
-${errors}
-`));
+    const logger = ai.getLogger();
+    logger(`\u274C Error Correction:
+${errors}`, {
+      tags: ["error"]
+    });
   }
 }
 
@@ -6932,7 +6994,6 @@ var validateJSONSchema = (schema) => {
 };
 
 // dsp/functions.ts
-var colorLog5 = new ColorLog();
 var AxFunctionError = class extends Error {
   constructor(fields) {
     super();
@@ -7101,12 +7162,11 @@ var processFunctions = async (ai, functionList, functionCalls, mem, sessionId, t
         );
         mem.addTag("error");
         if (ai.getOptions().debug) {
-          process.stdout.write(
-            colorLog5.red(`
-\u274C Function Error Correction:
-${result}
-`)
-          );
+          const logger = ai.getLogger();
+          logger(`\u274C Function Error Correction:
+${result}`, {
+            tags: ["error"]
+          });
         }
       } else {
         throw e;
@@ -7872,9 +7932,11 @@ var AxGen = class extends AxProgramWithSignature {
   values = {};
   excludeContentFromTrace = false;
   thoughtFieldName;
+  logger;
   constructor(signature, options) {
     super(signature, { description: options?.description });
     this.options = options;
+    this.logger = options?.logger;
     this.thoughtFieldName = options?.thoughtFieldName ?? "thought";
     const promptTemplateOptions = {
       functions: options?.functions,
@@ -8162,6 +8224,10 @@ Content: ${content}`
         xstate
       );
     }
+    if (ai.getOptions().debug) {
+      const logger = ai.getLogger();
+      logger("", { tags: ["responseEnd"] });
+    }
   }
   async processResponse({
     ai,
@@ -8289,7 +8355,8 @@ Content: ${result.content}`
             continue multiStepLoop;
           }
           if (debug) {
-            process.stdout.write("\n");
+            const logger = options.logger ?? this.logger ?? ai.getLogger();
+            logger("", { tags: ["responseEnd"] });
           }
           return;
         } catch (e) {
@@ -8480,10 +8547,14 @@ function processChildAgentFunction(childFunction, parentValues, parentInputKeys,
           ...pick(parentValues, injectionKeys)
         };
         if (options.debug && injectionKeys.length > 0) {
-          process.stdout.write(
-            `
-Function Params: ${JSON.stringify(updatedChildArgs, null, 2)}`
-          );
+          const ai = funcOptions?.ai;
+          if (ai) {
+            const logger = ai.getLogger();
+            logger(
+              `Function Params: ${JSON.stringify(updatedChildArgs, null, 2)}`,
+              { tags: ["functionArg"] }
+            );
+          }
         }
         return await originalFunc(updatedChildArgs, funcOptions);
       };
@@ -8592,18 +8663,18 @@ var AxAgent = class {
       }
       const debug = this.getDebug(ai, options);
       if (debug) {
-        process.stdout.write(`
---- Agent Engaged: ${this.name} ---
-`);
+        const logger = ai.getLogger();
+        logger(`\u{1F916} Agent ${this.name} starting...`, {
+          tags: ["assistantStart"]
+        });
       }
       const ret = await boundFunc(ai, values, {
         ...options,
         model
       });
       if (debug) {
-        process.stdout.write(`
---- Agent Done: ${this.name} ---
-`);
+        const logger = ai.getLogger();
+        logger(`\u{1F916} Agent ${this.name} completed.`, { tags: ["assistantEnd"] });
       }
       const sig = this.program.getSignature();
       const outFields = sig.getOutputFields();
@@ -9006,6 +9077,9 @@ var AxBalancer = class _AxBalancer {
   getOptions() {
     return this.currentService.getOptions();
   }
+  getLogger() {
+    return this.currentService.getLogger();
+  }
 };
 function validateModels2(services) {
   const serviceWithModel = services.find(
@@ -11441,6 +11515,11 @@ var AxMockAIService = class {
   getOptions() {
     return this.config.options ?? {};
   }
+  getLogger() {
+    return this.config.options?.logger ?? ((message) => {
+      process.stdout.write(message);
+    });
+  }
   updateMetrics(type) {
     const latency = this.config.latencyMs ?? 0;
     this.metrics.latency[type].samples.push(latency);
@@ -11463,7 +11542,7 @@ var AxMockAIService = class {
 };
 
 // dsp/classifier.ts
-var colorLog6 = new ColorLog();
+var colorLog4 = new ColorLog();
 var AxSimpleClassifierClass = class {
   name;
   context;
@@ -11525,7 +11604,7 @@ var AxSimpleClassifier = class {
     }
     if (this.debug) {
       console.log(
-        colorLog6.whiteBright(`query: ${text}`) + "\n" + colorLog6.greenBright(
+        colorLog4.whiteBright(`query: ${text}`) + "\n" + colorLog4.greenBright(
           JSON.stringify(m.map((m2) => `${m2.id}, ${m2.score}`))
         )
       );
@@ -13016,15 +13095,16 @@ function v4(options, buf, offset) {
 var v4_default = v4;
 
 // mcp/client.ts
-var colorLog7 = new ColorLog();
 var AxMCPClient = class {
   constructor(transport, options = {}) {
     this.transport = transport;
     this.options = options;
+    this.logger = options.logger ?? ((message) => console.log(message));
   }
   functions = [];
   activeRequests = /* @__PURE__ */ new Map();
   capabilities = {};
+  logger;
   async init() {
     if ("connect" in this.transport) {
       await this.transport.connect?.();
@@ -13083,11 +13163,13 @@ var AxMCPClient = class {
       };
     });
     if (this.options.debug) {
-      console.log(
-        colorLog7.yellow(`> Discovered ${this.functions.length} functions:`)
-      );
+      this.logger(`> Discovered ${this.functions.length} functions:`, {
+        tags: ["discovery"]
+      });
       for (const fn of this.functions) {
-        console.log(colorLog7.yellow(`  - ${fn.name}: ${fn.description}`));
+        this.logger(`  - ${fn.name}: ${fn.description}`, {
+          tags: ["discovery"]
+        });
       }
     }
   }
@@ -13130,11 +13212,10 @@ var AxMCPClient = class {
       params
     };
     if (this.options.debug) {
-      console.log(
-        colorLog7.blueBright(
-          `> Sending request ${requestId}:
-${JSON.stringify(request, null, 2)}`
-        )
+      this.logger(
+        `> Sending request ${requestId}:
+${JSON.stringify(request, null, 2)}`,
+        { tags: ["requestStart"] }
       );
     }
     const responsePromise = new Promise((resolve, reject) => {
@@ -13142,11 +13223,10 @@ ${JSON.stringify(request, null, 2)}`
       this.transport.send(request).then((res) => {
         this.activeRequests.delete(requestId);
         if (this.options.debug) {
-          console.log(
-            colorLog7.greenBright(
-              `> Received response for request ${requestId}:
-${JSON.stringify(res, null, 2)}`
-            )
+          this.logger(
+            `> Received response for request ${requestId}:
+${JSON.stringify(res, null, 2)}`,
+            { tags: ["responseContent"] }
           );
         }
         if (res !== null && typeof res === "object" && "error" in res) {
@@ -13176,9 +13256,9 @@ ${JSON.stringify(res, null, 2)}`
       params
     };
     if (this.options.debug) {
-      console.log(
-        "\u27A1\uFE0F Sending notification:",
-        JSON.stringify(notification, null, 2)
+      this.logger(
+        `\u27A1\uFE0F Sending notification: ${JSON.stringify(notification, null, 2)}`,
+        { tags: ["requestStart"] }
       );
     }
     await this.transport.sendNotification(notification);
@@ -13408,6 +13488,25 @@ var AxMultiServiceRouter = class {
   getOptions() {
     return this.options ?? {};
   }
+  /**
+   * Returns the logger from the last used service,
+   * or falls back to the first service if none has been used.
+   */
+  getLogger() {
+    let serviceInstance = this.lastUsedService;
+    if (!serviceInstance) {
+      const firstServiceEntry = this.services.values().next().value;
+      if (firstServiceEntry) {
+        serviceInstance = firstServiceEntry.service;
+      }
+    }
+    if (!serviceInstance) {
+      return (message) => {
+        process.stdout.write(message);
+      };
+    }
+    return serviceInstance.getLogger();
+  }
 };
 
 // prompts/rag.ts