@botpress/zai 2.1.19 → 2.2.0

package/dist/operations/group.js

@@ -0,0 +1,278 @@
+import { z } from "@bpinternal/zui";
+import { clamp } from "lodash-es";
+import pLimit from "p-limit";
+import { ZaiContext } from "../context";
+import { Response } from "../response";
+import { getTokenizer } from "../tokenizer";
+import { stringify } from "../utils";
+import { Zai } from "../zai";
+import { PROMPT_INPUT_BUFFER, PROMPT_OUTPUT_BUFFER } from "./constants";
+const _InitialGroup = z.object({
+  id: z.string().min(1).max(100),
+  label: z.string().min(1).max(250),
+  elements: z.array(z.any()).optional().default([])
+});
+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([])
+});
+const END = "\u25A0END\u25A0";
+const normalizeLabel = (label) => {
+  return label.trim().toLowerCase().replace(/^(group|new group|new)\s*[-:]\s*/i, "").replace(/^(group|new group|new)\s+/i, "").trim();
+};
+const group = async (input, _options, ctx) => {
+  ctx.controller.signal.throwIfAborted();
+  const options = _Options.parse(_options ?? {});
+  const tokenizer = await getTokenizer();
+  const model = await ctx.getModel();
+  if (input.length === 0) {
+    return [];
+  }
+  const groups = /* @__PURE__ */ new Map();
+  const groupElements = /* @__PURE__ */ new Map();
+  const elementGroups = /* @__PURE__ */ new Map();
+  const labelToGroupId = /* @__PURE__ */ new Map();
+  let groupIdCounter = 0;
+  options.initialGroups.forEach((ig) => {
+    const normalized = normalizeLabel(ig.label);
+    groups.set(ig.id, { id: ig.id, label: ig.label, normalizedLabel: normalized });
+    groupElements.set(ig.id, /* @__PURE__ */ new Set());
+    labelToGroupId.set(normalized, ig.id);
+  });
+  const elements = input.map((element, idx) => ({
+    element,
+    index: idx,
+    stringified: stringify(element, false)
+  }));
+  const TOKENS_TOTAL_MAX = model.input.maxTokens - PROMPT_INPUT_BUFFER - PROMPT_OUTPUT_BUFFER;
+  const TOKENS_INSTRUCTIONS_MAX = options.instructions ? clamp(tokenizer.count(options.instructions), 100, TOKENS_TOTAL_MAX * 0.2) : 0;
+  const TOKENS_AVAILABLE = TOKENS_TOTAL_MAX - TOKENS_INSTRUCTIONS_MAX;
+  const TOKENS_FOR_GROUPS_MAX = Math.floor(TOKENS_AVAILABLE * 0.4);
+  const TOKENS_FOR_ELEMENTS_MAX = Math.floor(TOKENS_AVAILABLE * 0.6);
+  const MAX_ELEMENTS_PER_CHUNK = 50;
+  const elementChunks = [];
+  let currentChunk = [];
+  let currentTokens = 0;
+  for (const elem of elements) {
+    const truncated = tokenizer.truncate(elem.stringified, options.tokensPerElement);
+    const elemTokens = tokenizer.count(truncated);
+    if ((currentTokens + elemTokens > TOKENS_FOR_ELEMENTS_MAX || currentChunk.length >= MAX_ELEMENTS_PER_CHUNK) && currentChunk.length > 0) {
+      elementChunks.push(currentChunk);
+      currentChunk = [];
+      currentTokens = 0;
+    }
+    currentChunk.push(elem.index);
+    currentTokens += elemTokens;
+  }
+  if (currentChunk.length > 0) {
+    elementChunks.push(currentChunk);
+  }
+  const getGroupChunks = () => {
+    const allGroupIds2 = Array.from(groups.keys());
+    if (allGroupIds2.length === 0) return [[]];
+    const chunks = [];
+    let currentChunk2 = [];
+    let currentTokens2 = 0;
+    for (const groupId of allGroupIds2) {
+      const group2 = groups.get(groupId);
+      const groupTokens = tokenizer.count(`${group2.label}`) + 10;
+      if (currentTokens2 + groupTokens > TOKENS_FOR_GROUPS_MAX && currentChunk2.length > 0) {
+        chunks.push(currentChunk2);
+        currentChunk2 = [];
+        currentTokens2 = 0;
+      }
+      currentChunk2.push(groupId);
+      currentTokens2 += groupTokens;
+    }
+    if (currentChunk2.length > 0) {
+      chunks.push(currentChunk2);
+    }
+    return chunks.length > 0 ? chunks : [[]];
+  };
+  const processChunk = async (elementIndices, groupIds) => {
+    const elementsText = elementIndices.map((idx, i) => {
+      const elem = elements[idx];
+      const truncated = tokenizer.truncate(elem.stringified, options.tokensPerElement);
+      return `\u25A0${i}: ${truncated}\u25A0`;
+    }).join("\n");
+    const groupsList = groupIds.map((gid) => groups.get(gid).label);
+    const groupsText = groupsList.length > 0 ? `**Existing Groups (prefer reusing these):**
+${groupsList.map((l) => `- ${l}`).join("\n")}
+
+` : "";
+    const systemPrompt = `You are grouping elements into cohesive groups.
+
+${options.instructions ? `**Instructions:** ${options.instructions}
+` : "**Instructions:** Group similar elements together."}
+
+**Important:**
+- Each element gets exactly ONE group label
+- Use EXACT SAME label for similar items (case-sensitive)
+- Create new descriptive labels when needed
+
+**Output Format:**
+One line per element:
+\u25A00:Group Label\u25A0
+\u25A01:Group Label\u25A0
+${END}`.trim();
+    const userPrompt = `${groupsText}**Elements (\u25A00 to \u25A0${elementIndices.length - 1}):**
+${elementsText}
+
+**Task:** For each element, output one line with its group label.
+${END}`.trim();
+    const { extracted } = await ctx.generateContent({
+      systemPrompt,
+      stopSequences: [END],
+      messages: [{ type: "text", role: "user", content: userPrompt }],
+      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 >= elementIndices.length) continue;
+          const label = (match[2] ?? "").trim();
+          if (!label) continue;
+          assignments.push({
+            elementIndex: elementIndices[idx],
+            label: label.slice(0, 250)
+          });
+        }
+        return assignments;
+      }
+    });
+    return extracted;
+  };
+  const elementLimit = pLimit(10);
+  const groupLimit = pLimit(10);
+  const allChunkResults = await Promise.all(
+    elementChunks.map(
+      (elementChunk) => elementLimit(async () => {
+        const groupChunks = getGroupChunks();
+        const allAssignments = await Promise.all(
+          groupChunks.map((groupChunk) => groupLimit(() => processChunk(elementChunk, groupChunk)))
+        );
+        return allAssignments.flat();
+      })
+    )
+  );
+  for (const assignments of allChunkResults) {
+    for (const { elementIndex, label } of assignments) {
+      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);
+      if (!elementGroups.has(elementIndex)) {
+        elementGroups.set(elementIndex, /* @__PURE__ */ new Set());
+      }
+      elementGroups.get(elementIndex).add(groupId);
+    }
+  }
+  const allGroupIds = Array.from(groups.keys());
+  if (allGroupIds.length > 0) {
+    const elementsNeedingReview = [];
+    for (const elem of elements) {
+      const seenGroups = elementGroups.get(elem.index) ?? /* @__PURE__ */ new Set();
+      const unseenCount = allGroupIds.filter((gid) => !seenGroups.has(gid)).length;
+      if (unseenCount > 0) {
+        elementsNeedingReview.push(elem.index);
+      }
+    }
+    if (elementsNeedingReview.length > 0) {
+      const reviewChunks = [];
+      let reviewChunk = [];
+      let reviewTokens = 0;
+      for (const elemIdx of elementsNeedingReview) {
+        const elem = elements[elemIdx];
+        const truncated = tokenizer.truncate(elem.stringified, options.tokensPerElement);
+        const elemTokens = tokenizer.count(truncated);
+        const shouldStartNewChunk = (reviewTokens + elemTokens > TOKENS_FOR_ELEMENTS_MAX || reviewChunk.length >= MAX_ELEMENTS_PER_CHUNK) && reviewChunk.length > 0;
+        if (shouldStartNewChunk) {
+          reviewChunks.push(reviewChunk);
+          reviewChunk = [];
+          reviewTokens = 0;
+        }
+        reviewChunk.push(elemIdx);
+        reviewTokens += elemTokens;
+      }
+      if (reviewChunk.length > 0) {
+        reviewChunks.push(reviewChunk);
+      }
+      const reviewResults = await Promise.all(
+        reviewChunks.map(
+          (chunk) => elementLimit(async () => {
+            const groupChunks = getGroupChunks();
+            const allAssignments = await Promise.all(
+              groupChunks.map((groupChunk) => groupLimit(() => processChunk(chunk, groupChunk)))
+            );
+            return allAssignments.flat();
+          })
+        )
+      );
+      const updateElementGroupAssignment = (elementIndex, label) => {
+        const normalized = normalizeLabel(label);
+        const groupId = labelToGroupId.get(normalized);
+        if (!groupId) return;
+        groupElements.get(groupId).add(elementIndex);
+        const elemGroups = elementGroups.get(elementIndex) ?? /* @__PURE__ */ new Set();
+        if (!elementGroups.has(elementIndex)) {
+          elementGroups.set(elementIndex, elemGroups);
+        }
+        elemGroups.add(groupId);
+      };
+      for (const assignments of reviewResults) {
+        for (const { elementIndex, label } of assignments) {
+          updateElementGroupAssignment(elementIndex, label);
+        }
+      }
+    }
+  }
+  for (const [elementIndex, groupSet] of elementGroups.entries()) {
+    if (groupSet.size > 1) {
+      const groupIds = Array.from(groupSet);
+      for (const gid of groupIds) {
+        groupElements.get(gid)?.delete(elementIndex);
+      }
+      const finalGroupId = groupIds[0];
+      groupElements.get(finalGroupId).add(elementIndex);
+    }
+  }
+  const result = [];
+  for (const [groupId, elementIndices] of groupElements.entries()) {
+    if (elementIndices.size > 0) {
+      const groupInfo = groups.get(groupId);
+      result.push({
+        id: groupInfo.id,
+        label: groupInfo.label,
+        elements: Array.from(elementIndices).map((idx) => elements[idx].element)
+      });
+    }
+  }
+  return result;
+};
+Zai.prototype.group = function(input, _options) {
+  const context = new ZaiContext({
+    client: this.client,
+    modelId: this.Model,
+    taskId: this.taskId,
+    taskType: "zai.group",
+    adapter: this.adapter
+  });
+  return new Response(context, group(input, _options, context), (result) => {
+    const merged = {};
+    result.forEach((group2) => {
+      if (!merged[group2.label]) {
+        merged[group2.label] = [];
+      }
+      merged[group2.label].push(...group2.elements);
+    });
+    return merged;
+  });
+};

package/dist/operations/label.js

@@ -1,5 +1,6 @@
 import { z } from "@bpinternal/zui";
 import { chunk, clamp } from "lodash-es";
+import pLimit from "p-limit";
 import { ZaiContext } from "../context";
 import { Response } from "../response";
 import { getTokenizer } from "../tokenizer";
@@ -87,9 +88,10 @@ const label = async (input, _labels, _options, ctx) => {
   );
   const inputAsString = stringify(input);
   if (tokenizer.count(inputAsString) > CHUNK_INPUT_MAX_TOKENS) {
+    const limit = pLimit(10);
     const tokens = tokenizer.split(inputAsString);
     const chunks = chunk(tokens, CHUNK_INPUT_MAX_TOKENS).map((x) => x.join(""));
-    const allLabels = await Promise.all(chunks.map((chunk2) => label(chunk2, _labels, _options, ctx)));
+    const allLabels = await Promise.all(chunks.map((chunk2) => limit(() => label(chunk2, _labels, _options, ctx))));
     return allLabels.reduce((acc, x) => {
       Object.keys(x).forEach((key) => {
         if (acc[key]?.value === true) {

package/dist/operations/summarize.js

@@ -1,5 +1,6 @@
 import { z } from "@bpinternal/zui";
 import { chunk } from "lodash-es";
+import pLimit from "p-limit";
 import { ZaiContext } from "../context";
 import { Response } from "../response";
 import { getTokenizer } from "../tokenizer";
@@ -54,8 +55,9 @@ ${newText}
   const useMergeSort = parts >= Math.pow(2, N);
   const chunkSize = Math.ceil(tokens.length / (parts * N));
   if (useMergeSort) {
+    const limit = pLimit(10);
     const chunks = chunk(tokens, chunkSize).map((x) => x.join(""));
-    const allSummaries = (await Promise.allSettled(chunks.map((chunk2) => summarize(chunk2, options, ctx)))).filter((x) => x.status === "fulfilled").map((x) => x.value);
+    const allSummaries = (await Promise.allSettled(chunks.map((chunk2) => limit(() => summarize(chunk2, options, ctx))))).filter((x) => x.status === "fulfilled").map((x) => x.value);
     return summarize(allSummaries.join("\n\n============\n\n"), options, ctx);
   }
   const summaries = [];