@botpress/zai 2.1.20 → 2.2.0

package/README.md

@@ -104,7 +104,32 @@ const techCompanies = await zai.filter(companies, 'are technology companies')
 const recentPosts = await zai.filter(posts, 'were published this week')
 ```
 
-### 6. Text - Generate content
+### 6. Group - Organize items into categories
+
+```typescript
+// Group items automatically
+const grouped = await zai.group(tasks, {
+  instructions: 'Group by priority level',
+})
+// Result: { 'High Priority': [...], 'Medium Priority': [...], 'Low Priority': [...] }
+
+// Group with initial categories
+const categorized = await zai.group(emails, {
+  instructions: 'Group by topic',
+  initialGroups: [
+    { id: 'work', label: 'Work' },
+    { id: 'personal', label: 'Personal' },
+  ],
+})
+
+// Group large datasets efficiently
+const organized = await zai.group(largeArray, {
+  instructions: 'Group by date',
+  chunkLength: 8000, // Process in chunks for better performance
+})
+```
+
+### 7. Text - Generate content
 
 ```typescript
 const blogPost = await zai.text('Write about the future of AI', {
@@ -113,7 +138,7 @@ const blogPost = await zai.text('Write about the future of AI', {
 })
 ```
 
-### 7. Summarize - Create summaries
+### 8. Summarize - Create summaries
 
 ```typescript
 // Simple summary
@@ -237,6 +262,7 @@ setTimeout(() => controller.abort(), 5000)
 - `.label(content, criteria, options?)` - Apply multiple labels
 - `.rewrite(content, instruction, options?)` - Transform text
 - `.filter(items, condition, options?)` - Filter array items
+- `.group(items, options?)` - Organize items into categories
 - `.text(prompt, options?)` - Generate text
 - `.summarize(content, options?)` - Create summary
 

package/dist/index.d.ts

@@ -166,14 +166,14 @@ declare class Response<T = any, S = T> implements PromiseLike<S> {
     }>;
 }
 
-type Options$6 = {
+type Options$7 = {
     /** The maximum number of tokens to generate */
     length?: number;
 };
 declare module '@botpress/zai' {
     interface Zai {
         /** Generates a text of the desired length according to the prompt */
-        text(prompt: string, options?: Options$6): Response<string>;
+        text(prompt: string, options?: Options$7): Response<string>;
     }
 }
 
@@ -182,7 +182,7 @@ type Example$3 = {
     output: string;
     instructions?: string;
 };
-type Options$5 = {
+type Options$6 = {
     /** Examples to guide the rewriting */
     examples?: Array<Example$3>;
     /** The maximum number of tokens to generate */
@@ -191,11 +191,11 @@ type Options$5 = {
 declare module '@botpress/zai' {
     interface Zai {
         /** Rewrites a string according to match the prompt */
-        rewrite(original: string, prompt: string, options?: Options$5): Response<string>;
+        rewrite(original: string, prompt: string, options?: Options$6): Response<string>;
     }
 }
 
-type Options$4 = {
+type Options$5 = {
     /** What should the text be summarized to? */
     prompt?: string;
     /** How to format the example text */
@@ -215,7 +215,7 @@ type Options$4 = {
 declare module '@botpress/zai' {
     interface Zai {
         /** Summarizes a text of any length to a summary of the desired length */
-        summarize(original: string, options?: Options$4): Response<string>;
+        summarize(original: string, options?: Options$5): Response<string>;
     }
 }
 
@@ -225,14 +225,14 @@ type Example$2 = {
     reason?: string;
     condition?: string;
 };
-type Options$3 = {
+type Options$4 = {
     /** Examples to check the condition against */
     examples?: Array<Example$2>;
 };
 declare module '@botpress/zai' {
     interface Zai {
         /** Checks wether a condition is true or not */
-        check(input: unknown, condition: string, options?: Options$3): Response<{
+        check(input: unknown, condition: string, options?: Options$4): Response<{
             /** Whether the condition is true or not */
             value: boolean;
             /** The explanation of the decision */
@@ -246,7 +246,7 @@ type Example$1 = {
     filter: boolean;
     reason?: string;
 };
-type Options$2 = {
+type Options$3 = {
     /** The maximum number of tokens per item */
     tokensPerItem?: number;
     /** Examples to filter the condition against */
@@ -255,11 +255,11 @@ type Options$2 = {
 declare module '@botpress/zai' {
     interface Zai {
         /** Filters elements of an array against a condition */
-        filter<T>(input: Array<T>, condition: string, options?: Options$2): Response<Array<T>>;
+        filter<T>(input: Array<T>, condition: string, options?: Options$3): Response<Array<T>>;
     }
 }
 
-type Options$1 = {
+type Options$2 = {
     /** Instructions to guide the user on how to extract the data */
     instructions?: string;
     /** The maximum number of tokens per chunk */
@@ -274,7 +274,7 @@ type OfType<O, T extends __Z = __Z<O>> = T extends __Z<O> ? T : never;
 declare module '@botpress/zai' {
     interface Zai {
         /** Extracts one or many elements from an arbitrary input */
-        extract<S extends OfType<any>>(input: unknown, schema: S, options?: Options$1): Response<S['_output']>;
+        extract<S extends OfType<any>>(input: unknown, schema: S, options?: Options$2): Response<S['_output']>;
     }
 }
 
@@ -293,7 +293,7 @@ type Example<T extends string> = {
         explanation?: string;
     }>>;
 };
-type Options<T extends string> = {
+type Options$1<T extends string> = {
     /** Examples to help the user make a decision */
     examples?: Array<Example<T>>;
     /** Instructions to guide the user on how to extract the data */
@@ -305,7 +305,7 @@ type Labels<T extends string> = Record<T, string>;
 declare module '@botpress/zai' {
     interface Zai {
         /** Tags the provided input with a list of predefined labels */
-        label<T extends string>(input: unknown, labels: Labels<T>, options?: Options<T>): Response<{
+        label<T extends string>(input: unknown, labels: Labels<T>, options?: Options$1<T>): Response<{
             [K in T]: {
                 explanation: string;
                 value: boolean;
@@ -317,4 +317,26 @@ declare module '@botpress/zai' {
     }
 }
 
+type Group<T> = {
+    id: string;
+    label: string;
+    elements: T[];
+};
+type InitialGroup = {
+    id: string;
+    label: string;
+    elements?: unknown[];
+};
+type Options = {
+    instructions?: string;
+    tokensPerElement?: number;
+    chunkLength?: number;
+    initialGroups?: Array<InitialGroup>;
+};
+declare module '@botpress/zai' {
+    interface Zai {
+        group<T>(input: Array<T>, options?: Options): Response<Array<Group<T>>, Record<string, T[]>>;
+    }
+}
+
 export { Zai };

package/dist/index.js

@@ -6,4 +6,5 @@ import "./operations/check";
 import "./operations/filter";
 import "./operations/extract";
 import "./operations/label";
+import "./operations/group";
 export { Zai };

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;
+  });
+};