@botpress/zai 2.5.17 → 2.6.0

package/dist/context.js

@@ -1,4 +1,5 @@
 import { EventEmitter } from "./emitter";
+import { fastHash } from "./utils";
 export class ZaiContext {
   _startedAt = Date.now();
   _inputCost = 0;
@@ -15,8 +16,10 @@ export class ZaiContext {
   adapter;
   source;
   _eventEmitter;
+  _memoizer;
   controller = new AbortController();
   _client;
+  static _noopMemoizer = { run: (_id, fn) => fn() };
   constructor(props) {
     this._client = props.client.clone();
     this.taskId = props.taskId;
@@ -24,6 +27,7 @@ export class ZaiContext {
     this.adapter = props.adapter;
     this.source = props.source;
     this.taskType = props.taskType;
+    this._memoizer = props.memoizer ?? ZaiContext._noopMemoizer;
     this._eventEmitter = new EventEmitter();
     this._client.on("request", () => {
       this._totalRequests++;
@@ -57,6 +61,16 @@ export class ZaiContext {
     this._eventEmitter.clear();
   }
   async generateContent(props) {
+    const memoKey = `zai:memo:${this.taskType}:${this.taskId || "default"}:${fastHash(
+      JSON.stringify({
+        s: props.systemPrompt,
+        m: props.messages?.map((m) => "content" in m ? m.content : ""),
+        st: props.stopSequences
+      })
+    )}`;
+    return this._memoizer.run(memoKey, () => this._generateContentInner(props));
+  }
+  async _generateContentInner(props) {
     const maxRetries = Math.max(props.maxRetries ?? 3, 0);
     const transform = props.transform;
     let lastError = null;

package/dist/index.d.ts

@@ -41,6 +41,16 @@ declare abstract class Adapter {
     abstract saveExample<TInput, TOutput>(props: SaveExampleProps<TInput, TOutput>): Promise<void>;
 }
 
+/**
+ * A memoizer that caches the result of async operations by a unique key.
+ *
+ * When used with the Botpress ADK workflow `step` function, this enables
+ * Zai operations to resume where they left off if a workflow is interrupted.
+ *
+ */
+type Memoizer = {
+    run: <T>(id: string, fn: () => Promise<T>) => Promise<T>;
+};
 /**
  * Active learning configuration for improving AI operations over time.
  *
@@ -98,6 +108,16 @@ type ZaiConfig = {
     activeLearning?: ActiveLearning;
     /** Namespace for organizing tasks (default: 'zai') */
     namespace?: string;
+    /**
+     * Memoizer (or factory returning one) for caching cognitive call results.
+     *
+     * When provided, all LLM calls are wrapped in the memoizer, allowing results
+     * to be cached and replayed. This is useful for resuming workflow runs where
+     * Zai operations have already completed their cognitive calls.
+     *
+     * If a factory function is provided, it is called once per Zai operation invocation.
+     */
+    memoize?: Memoizer | (() => Memoizer);
 };
 /**
  * Zai - A type-safe LLM utility library for production-ready AI operations.
@@ -171,6 +191,7 @@ declare class Zai {
     protected namespace: string;
     protected adapter: Adapter;
     protected activeLearning: ActiveLearning;
+    protected _memoize?: Memoizer | (() => Memoizer);
     /**
      * Creates a new Zai instance with the specified configuration.
      *
@@ -195,6 +216,8 @@ declare class Zai {
     constructor(config: ZaiConfig);
     /** @internal */
     protected callModel(props: Parameters<Cognitive['generateContent']>[0]): ReturnType<Cognitive['generateContent']>;
+    /** @internal */
+    protected _resolveMemoizer(): Memoizer | undefined;
     protected getTokenizer(): Promise<TextTokenizer>;
     protected fetchModelDetails(): Promise<void>;
     protected get taskId(): string;
@@ -299,6 +322,7 @@ type ZaiContextProps = {
     modelId: string;
     adapter?: Adapter;
     source?: GenerateContentInput['meta'];
+    memoizer?: Memoizer;
 };
 /**
  * Usage statistics tracking tokens, cost, and request metrics for an operation.
@@ -370,8 +394,10 @@ declare class ZaiContext {
     adapter?: Adapter;
     source?: GenerateContentInput['meta'];
     private _eventEmitter;
+    private _memoizer;
     controller: AbortController;
     private _client;
+    private static _noopMemoizer;
     constructor(props: ZaiContextProps);
     getModel(): Promise<Model>;
     on<K extends keyof ContextEvents>(type: K, listener: (event: ContextEvents[K]) => void): this;
@@ -382,6 +408,7 @@ declare class ZaiContext {
         text: string | undefined;
         extracted: Out;
     }>;
+    private _generateContentInner;
     get elapsedTime(): number;
     get usage(): Usage;
 }
@@ -1390,6 +1417,8 @@ type Options$4 = {
     tokensPerElement?: number;
     chunkLength?: number;
     initialGroups?: Array<InitialGroup>;
+    maxGroups?: number;
+    minElements?: number;
 };
 declare module '@botpress/zai' {
     interface Zai {
@@ -1402,6 +1431,8 @@ declare module '@botpress/zai' {
          *
          * @param input - Array of items to group
          * @param options - Configuration for grouping behavior, instructions, and initial categories
+         * @param options.maxGroups - Maximum number of groups allowed (minimum 2). When set, groups are merged at the end until within limit.
+         * @param options.minElements - Minimum elements per group (minimum 1). Groups below this threshold have their elements redistributed via AI.
          * @returns Response with groups array (simplified to Record<groupLabel, items[]>)
          *
          * @example Automatic grouping
@@ -1543,6 +1574,18 @@ declare module '@botpress/zai' {
          * })
          * ```
          */
+        /**
+         * @example Limiting number of groups
+         * ```typescript
+         * const items = ['apple', 'banana', 'carrot', 'chicken', 'rice', 'bread', 'salmon', 'milk']
+         *
+         * const groups = await zai.group(items, {
+         *   instructions: 'Group by food type',
+         *   maxGroups: 3 // At most 3 groups — smallest groups get merged if exceeded
+         * })
+         * // Guarantees no more than 3 groups in the result
+         * ```
+         */
         group<T>(input: Array<T>, options?: Options$4): Response<Array<Group<T>>, Record<string, T[]>>;
     }
 }
@@ -2127,4 +2170,4 @@ declare module '@botpress/zai' {
     }
 }
 
-export { Zai };
+export { type Memoizer, Zai };

package/dist/operations/answer.js

@@ -373,7 +373,8 @@ Zai.prototype.answer = function(documents, question, _options) {
     modelId: this.Model,
     taskId: this.taskId,
     taskType: "zai.answer",
-    adapter: this.adapter
+    adapter: this.adapter,
+    memoizer: this._resolveMemoizer()
   });
   return new Response(
     context,

package/dist/operations/check.js

@@ -181,7 +181,8 @@ Zai.prototype.check = function(input, condition, _options) {
     modelId: this.Model,
     taskId: this.taskId,
     taskType: "zai.check",
-    adapter: this.adapter
+    adapter: this.adapter,
+    memoizer: this._resolveMemoizer()
   });
   return new Response(context, check(input, condition, options, context), (result) => result.value);
 };

package/dist/operations/extract.js

@@ -313,7 +313,8 @@ Zai.prototype.extract = function(input, schema, _options) {
     modelId: this.Model,
     taskId: this.taskId,
     taskType: "zai.extract",
-    adapter: this.adapter
+    adapter: this.adapter,
+    memoizer: this._resolveMemoizer()
   });
   return new Response(context, extract(input, schema, _options, context), (result) => result);
 };

package/dist/operations/filter.js

@@ -202,7 +202,8 @@ Zai.prototype.filter = function(input, condition, _options) {
     modelId: this.Model,
     taskId: this.taskId,
     taskType: "zai.filter",
-    adapter: this.adapter
+    adapter: this.adapter,
+    memoizer: this._resolveMemoizer()
   });
   return new Response(context, filter(input, condition, _options, context), (result) => result);
 };

package/dist/operations/group.js

@@ -16,7 +16,9 @@ const _Options = z.object({
   instructions: z.string().optional(),
   tokensPerElement: z.number().min(1).max(1e5).optional().default(250),
   chunkLength: z.number().min(100).max(1e5).optional().default(16e3),
-  initialGroups: z.array(_InitialGroup).optional().default([])
+  initialGroups: z.array(_InitialGroup).optional().default([]),
+  maxGroups: z.number().min(2).optional(),
+  minElements: z.number().min(1).optional()
 });
 const END = "\u25A0END\u25A0";
 const normalizeLabel = (label) => {
@@ -301,6 +303,191 @@ ${END}`.trim();
       groupElements.get(finalGroupId).add(elementIndex);
     }
   }
+  if (options.maxGroups !== void 0) {
+    const nonEmptyGroupIds = () => Array.from(groupElements.entries()).filter(([, s]) => s.size > 0).map(([id]) => id);
+    let currentIds = nonEmptyGroupIds();
+    if (currentIds.length > options.maxGroups) {
+      const groupSummaries = currentIds.map((gid, idx) => {
+        const info = groups.get(gid);
+        const elemIndices = Array.from(groupElements.get(gid));
+        const sampleElements = elemIndices.slice(0, 3).map((i) => tokenizer.truncate(elements[i].stringified, 60)).join(", ");
+        return `\u25A0${idx}:${info.label} (${elemIndices.length} elements, e.g. ${sampleElements})\u25A0`;
+      });
+      const mergeSystemPrompt = `You are consolidating groups into fewer, broader categories.
+
+${options.instructions ? `**Original instructions:** ${options.instructions}
+` : ""}
+**Task:** Merge ${currentIds.length} groups down to at most ${options.maxGroups} groups.
+Combine the most semantically related groups together. Give each merged group a new descriptive label.
+
+**Output Format:**
+For each input group (\u25A00 to \u25A0${currentIds.length - 1}), output which target label it maps to:
+\u25A00:Merged Label\u25A0
+\u25A01:Merged Label\u25A0
+${END}
+
+Use the EXACT SAME label for groups that should be merged together.`.trim();
+      const mergeUserPrompt = `**Current groups:**
+${groupSummaries.join("\n")}
+
+Merge into at most ${options.maxGroups} groups.
+${END}`.trim();
+      const { extracted: mergeAssignments } = await ctx.generateContent({
+        systemPrompt: mergeSystemPrompt,
+        stopSequences: [END],
+        messages: [{ type: "text", role: "user", content: mergeUserPrompt }],
+        transform: (text) => {
+          const assignments = [];
+          const regex = /■(\d+):([^■]+)■/g;
+          let match;
+          while ((match = regex.exec(text)) !== null) {
+            const idx = parseInt(match[1] ?? "", 10);
+            if (isNaN(idx) || idx < 0 || idx >= currentIds.length) continue;
+            const label = (match[2] ?? "").trim();
+            if (!label) continue;
+            assignments.push({ sourceIdx: idx, label: label.slice(0, 250) });
+          }
+          return assignments;
+        }
+      });
+      const mergeMap = /* @__PURE__ */ new Map();
+      for (const { sourceIdx, label } of mergeAssignments) {
+        const sourceGid = currentIds[sourceIdx];
+        if (!sourceGid) continue;
+        const normalized = normalizeLabel(label);
+        if (!mergeMap.has(normalized)) {
+          mergeMap.set(normalized, { label, sourceGroupIds: [] });
+        }
+        mergeMap.get(normalized).sourceGroupIds.push(sourceGid);
+      }
+      for (const [, { label, sourceGroupIds }] of mergeMap) {
+        if (sourceGroupIds.length <= 1) continue;
+        const targetGid = sourceGroupIds[0];
+        const targetSet = groupElements.get(targetGid);
+        const targetInfo = groups.get(targetGid);
+        targetInfo.label = label;
+        targetInfo.normalizedLabel = normalizeLabel(label);
+        for (let i = 1; i < sourceGroupIds.length; i++) {
+          const sourceGid = sourceGroupIds[i];
+          const sourceSet = groupElements.get(sourceGid);
+          sourceSet.forEach((elemIdx) => targetSet.add(elemIdx));
+          sourceSet.clear();
+        }
+      }
+      currentIds = nonEmptyGroupIds();
+      while (currentIds.length > options.maxGroups) {
+        currentIds.sort((a, b) => groupElements.get(a).size - groupElements.get(b).size);
+        const sourceSet = groupElements.get(currentIds[0]);
+        const targetSet = groupElements.get(currentIds[1]);
+        for (const elemIdx of sourceSet) {
+          targetSet.add(elemIdx);
+        }
+        sourceSet.clear();
+        currentIds = nonEmptyGroupIds();
+      }
+    }
+  }
+  if (options.minElements !== void 0 && options.minElements > 1) {
+    const getNonEmptyGroupIds = () => Array.from(groupElements.entries()).filter(([, s]) => s.size > 0).map(([id]) => id);
+    const orphanIndices = [];
+    for (const gid of getNonEmptyGroupIds()) {
+      const elemSet = groupElements.get(gid);
+      if (elemSet.size > 0 && elemSet.size < options.minElements) {
+        for (const idx of elemSet) {
+          orphanIndices.push(idx);
+        }
+        elemSet.clear();
+      }
+    }
+    if (orphanIndices.length > 0) {
+      const validGroupIds = getNonEmptyGroupIds();
+      const orphanChunks = [];
+      let currentOrphanChunk = [];
+      let currentOrphanTokens = 0;
+      for (const elemIdx of orphanIndices) {
+        const elem = elements[elemIdx];
+        const truncated = tokenizer.truncate(elem.stringified, options.tokensPerElement);
+        const elemTokens = tokenizer.count(truncated);
+        if ((currentOrphanTokens + elemTokens > TOKENS_FOR_ELEMENTS_MAX || currentOrphanChunk.length >= MAX_ELEMENTS_PER_CHUNK) && currentOrphanChunk.length > 0) {
+          orphanChunks.push(currentOrphanChunk);
+          currentOrphanChunk = [];
+          currentOrphanTokens = 0;
+        }
+        currentOrphanChunk.push(elemIdx);
+        currentOrphanTokens += elemTokens;
+      }
+      if (currentOrphanChunk.length > 0) {
+        orphanChunks.push(currentOrphanChunk);
+      }
+      const orphanResults = await Promise.all(
+        orphanChunks.map(
+          (chunk) => elementLimit(async () => {
+            const groupChunksForOrphans = validGroupIds.length > 0 ? getGroupChunks() : [[]];
+            const allAssignments = await Promise.all(
+              groupChunksForOrphans.filter((gc) => gc.length === 0 || gc.some((gid) => validGroupIds.includes(gid))).map((groupChunk) => {
+                const filteredGroupChunk = groupChunk.filter((gid) => validGroupIds.includes(gid));
+                return groupLimit(() => processChunk(chunk, filteredGroupChunk));
+              })
+            );
+            return allAssignments.flat();
+          })
+        )
+      );
+      const flatAssignments = orphanResults.flat();
+      for (const { elementIndex, label } of flatAssignments) {
+        const normalized = normalizeLabel(label);
+        let groupId = labelToGroupId.get(normalized);
+        if (!groupId) {
+          groupId = `group_${groupIdCounter++}`;
+          groups.set(groupId, { id: groupId, label, normalizedLabel: normalized });
+          groupElements.set(groupId, /* @__PURE__ */ new Set());
+          labelToGroupId.set(normalized, groupId);
+        }
+        groupElements.get(groupId).add(elementIndex);
+      }
+      const isAssigned = (idx) => {
+        for (const [, elemSet] of groupElements) {
+          if (elemSet.has(idx)) return true;
+        }
+        return false;
+      };
+      const unassigned = orphanIndices.filter((idx) => !isAssigned(idx));
+      const placeIntoLargest = (indices) => {
+        const allNonEmpty = getNonEmptyGroupIds();
+        if (allNonEmpty.length === 0) return;
+        const largestGid = allNonEmpty.reduce(
+          (a, b) => groupElements.get(a).size >= groupElements.get(b).size ? a : b
+        );
+        for (const idx of indices) {
+          groupElements.get(largestGid).add(idx);
+        }
+      };
+      if (unassigned.length > 0) {
+        placeIntoLargest(unassigned);
+      }
+      const mergeUndersizedGroups = () => {
+        const allNonEmpty = getNonEmptyGroupIds();
+        if (allNonEmpty.length <= 1) return false;
+        const largestGid = allNonEmpty.reduce(
+          (a, b) => groupElements.get(a).size >= groupElements.get(b).size ? a : b
+        );
+        const targetSet = groupElements.get(largestGid);
+        let merged = false;
+        for (const gid of allNonEmpty) {
+          if (gid === largestGid) continue;
+          const elemSet = groupElements.get(gid);
+          if (elemSet.size > 0 && elemSet.size < options.minElements) {
+            elemSet.forEach((idx) => targetSet.add(idx));
+            elemSet.clear();
+            merged = true;
+          }
+        }
+        return merged;
+      };
+      while (mergeUndersizedGroups()) {
+      }
+    }
+  }
   const result = [];
   for (const [groupId, elementIndices] of groupElements.entries()) {
     if (elementIndices.size > 0) {
@@ -354,7 +541,8 @@ Zai.prototype.group = function(input, _options) {
     modelId: this.Model,
     taskId: this.taskId,
     taskType: "zai.group",
-    adapter: this.adapter
+    adapter: this.adapter,
+    memoizer: this._resolveMemoizer()
   });
   return new Response(context, group(input, _options, context), (result) => {
     const merged = {};

package/dist/operations/label.js

@@ -276,7 +276,8 @@ Zai.prototype.label = function(input, labels, _options) {
     modelId: this.Model,
     taskId: this.taskId,
     taskType: "zai.label",
-    adapter: this.adapter
+    adapter: this.adapter,
+    memoizer: this._resolveMemoizer()
   });
   return new Response(
     context,

package/dist/operations/patch.js

@@ -392,7 +392,8 @@ Zai.prototype.patch = function(files, instructions, _options) {
     modelId: this.Model,
     taskId: this.taskId,
     taskType: "zai.patch",
-    adapter: this.adapter
+    adapter: this.adapter,
+    memoizer: this._resolveMemoizer()
   });
   return new Response(context, patch(files, instructions, _options, context), (result) => result);
 };

package/dist/operations/rate.js

@@ -335,7 +335,8 @@ Zai.prototype.rate = function(input, instructions, _options) {
     modelId: this.Model,
     taskId: this.taskId,
     taskType: "zai.rate",
-    adapter: this.adapter
+    adapter: this.adapter,
+    memoizer: this._resolveMemoizer()
   });
   return new Response(
     context,

package/dist/operations/rewrite.js

@@ -136,7 +136,8 @@ Zai.prototype.rewrite = function(original, prompt, _options) {
     modelId: this.Model,
     taskId: this.taskId,
     taskType: "zai.rewrite",
-    adapter: this.adapter
+    adapter: this.adapter,
+    memoizer: this._resolveMemoizer()
   });
   return new Response(context, rewrite(original, prompt, _options, context), (result) => result);
 };

package/dist/operations/sort.js

@@ -511,7 +511,8 @@ Zai.prototype.sort = function(input, instructions, _options) {
     modelId: this.Model,
     taskId: this.taskId,
     taskType: "zai.sort",
-    adapter: this.adapter
+    adapter: this.adapter,
+    memoizer: this._resolveMemoizer()
   });
   return new Response(
     context,

package/dist/operations/summarize.js

@@ -148,7 +148,8 @@ Zai.prototype.summarize = function(original, _options) {
     modelId: this.Model,
     taskId: this.taskId,
     taskType: "summarize",
-    adapter: this.adapter
+    adapter: this.adapter,
+    memoizer: this._resolveMemoizer()
   });
   return new Response(context, summarize(original, options, context), (value) => value);
 };

package/dist/operations/text.js

@@ -60,7 +60,8 @@ Zai.prototype.text = function(prompt, _options) {
     modelId: this.Model,
     taskId: this.taskId,
     taskType: "zai.text",
-    adapter: this.adapter
+    adapter: this.adapter,
+    memoizer: this._resolveMemoizer()
   });
   return new Response(context, text(prompt, _options, context), (result) => result);
 };

package/dist/zai.js

@@ -47,6 +47,7 @@ export class Zai {
   namespace;
   adapter;
   activeLearning;
+  _memoize;
   /**
    * Creates a new Zai instance with the specified configuration.
    *
@@ -80,6 +81,7 @@ export class Zai {
       client: this.client.client,
       tableName: parsed.activeLearning.tableName
     }) : new MemoryAdapter([]);
+    this._memoize = config.memoize;
   }
   /** @internal */
   async callModel(props) {
@@ -90,6 +92,13 @@ export class Zai {
       userId: this._userId
     });
   }
+  /** @internal */
+  _resolveMemoizer() {
+    if (!this._memoize) {
+      return void 0;
+    }
+    return typeof this._memoize === "function" ? this._memoize() : this._memoize;
+  }
   async getTokenizer() {
     Zai.tokenizer ??= await (async () => {
       while (!getWasmTokenizer) {