@botpress/zai 2.1.20 → 2.3.0

package/dist/operations/sort.js

@@ -0,0 +1,450 @@
+import { z } from "@bpinternal/zui";
+import pLimit from "p-limit";
+import { ZaiContext } from "../context";
+import { Response } from "../response";
+import { getTokenizer } from "../tokenizer";
+import { fastHash, stringify } from "../utils";
+import { Zai } from "../zai";
+import { PROMPT_INPUT_BUFFER, PROMPT_OUTPUT_BUFFER } from "./constants";
+const _Options = z.object({
+  tokensPerItem: z.number().min(1).max(1e5).optional().describe("The maximum number of tokens per item").default(250)
+});
+const END = "\u25A0END\u25A0";
+const sort = async (input, instructions, _options, ctx) => {
+  ctx.controller.signal.throwIfAborted();
+  const options = _Options.parse(_options ?? {});
+  const tokenizer = await getTokenizer();
+  const model = await ctx.getModel();
+  const taskId = ctx.taskId;
+  const taskType = "zai.sort";
+  if (input.length === 0) {
+    return [];
+  }
+  if (input.length === 1) {
+    return input;
+  }
+  const TOKENS_TOTAL_MAX = model.input.maxTokens - PROMPT_INPUT_BUFFER - PROMPT_OUTPUT_BUFFER;
+  const sampleSize = Math.min(5, input.length);
+  const sampleItems = input.slice(0, sampleSize);
+  const sampleItemsText = sampleItems.map((item, idx) => `\u25A0${idx}: ${stringify(item, false)}`).join("\n");
+  const generateCriteriaPrompt = `Analyze this sorting instruction: "${instructions}"
+
+Sample items to be sorted:
+${sampleItemsText}
+
+Create 1-3 sorting criteria with ordered label arrays (3-10 labels each).
+
+**CRITICAL RULES**:
+1. Labels are single words, lowercase, no spaces, use underscores
+2. Labels are ordered from FIRST to LAST in sorted result
+3. If instruction says "from X to Y": first label represents X, last label represents Y
+4. If instruction says "prioritize" or "highest/lowest priority":
+   - First label = HIGHEST priority (top of todo list)
+   - Last label = LOWEST priority (bottom of todo list)
+
+Examples:
+
+"from slowest to fastest" \u2192 first=slowest, last=fastest
+\u25A0speed\u25A0
+very_slow;slow;medium;fast;very_fast
+\u25A0END\u25A0
+
+"from most dangerous to least dangerous" \u2192 first=most dangerous, last=least dangerous
+\u25A0danger\u25A0
+extremely_dangerous;very_dangerous;dangerous;moderate;slightly_dangerous;harmless
+\u25A0END\u25A0
+
+"from least urgent (spam) to most urgent (bills)" \u2192 first=spam, last=bills
+\u25A0urgency\u25A0
+spam;promotional;normal;important;urgent;critical
+\u25A0END\u25A0
+
+"prioritize: highest priority=open old tickets; lowest priority=closed" \u2192 first=high priority, last=low priority
+\u25A0status\u25A0
+open_old;open_recent;closed
+\u25A0age\u25A0
+oldest;old;recent;new
+\u25A0END\u25A0
+
+Output format:
+\u25A0criterion_name\u25A0
+label1;label2;label3;label4
+\u25A0END\u25A0
+
+Use 3-10 labels per criterion. Labels should be intuitive and match the domain.
+Keep criterion names short (1-2 words, lowercase, underscores).
+`;
+  const { extracted: sortingCriteria } = await ctx.generateContent({
+    systemPrompt: `You are creating sorting criteria with ordered label arrays.
+
+CRITICAL: Output ordered labels from FIRST to LAST position in sorted result.
+- Labels are single words, lowercase, underscores only
+- 3-10 labels per criterion
+- Order matters: first label = appears first, last label = appears last`,
+    messages: [
+      {
+        type: "text",
+        role: "user",
+        content: generateCriteriaPrompt
+      }
+    ],
+    transform: (text) => {
+      const criteria = {};
+      const criterionRegex = /■([^■]+)■\s*([^\n■]+)/g;
+      let match;
+      while ((match = criterionRegex.exec(text)) !== null) {
+        const name = (match[1] ?? "").trim().toLowerCase();
+        const labelsStr = (match[2] ?? "").trim();
+        if (!name || name === "end") continue;
+        const labels = labelsStr.split(";").map((l) => l.trim().toLowerCase().replace(/\s+/g, "_")).filter((l) => l.length > 0 && l.length < 50);
+        if (labels.length >= 3 && labels.length <= 10) {
+          criteria[name] = {
+            description: `${labels.length} ordered labels`,
+            labels
+          };
+        }
+      }
+      if (Object.keys(criteria).length === 0) {
+        throw new Error(`Failed to parse sorting criteria. LLM output: ${text.slice(0, 500)}`);
+      }
+      return criteria;
+    }
+  });
+  const criteriaKeys = Object.keys(sortingCriteria);
+  if (criteriaKeys.length === 0) {
+    throw new Error("No sorting criteria generated");
+  }
+  const TOKENS_CRITERIA_MAX = Math.floor(TOKENS_TOTAL_MAX * 0.2);
+  const TOKENS_ITEMS_MAX = TOKENS_TOTAL_MAX - TOKENS_CRITERIA_MAX;
+  const MAX_ITEMS_PER_CHUNK = 50;
+  const elements = input.map((element, idx) => ({
+    element,
+    index: idx,
+    stringified: stringify(element, false)
+  }));
+  const chunks = [];
+  let currentChunk = [];
+  let currentTokens = 0;
+  for (const elem of elements) {
+    const truncated = tokenizer.truncate(elem.stringified, options.tokensPerItem);
+    const elemTokens = tokenizer.count(truncated);
+    if ((currentTokens + elemTokens > TOKENS_ITEMS_MAX || currentChunk.length >= MAX_ITEMS_PER_CHUNK) && currentChunk.length > 0) {
+      chunks.push(currentChunk);
+      currentChunk = [];
+      currentTokens = 0;
+    }
+    currentChunk.push(elem);
+    currentTokens += elemTokens;
+  }
+  if (currentChunk.length > 0) {
+    chunks.push(currentChunk);
+  }
+  const scoreChunk = async (chunk) => {
+    ctx.controller.signal.throwIfAborted();
+    const chunkSize = chunk.length;
+    const chunkInputStr = JSON.stringify(chunk.map((c) => c.element));
+    const examples = taskId && ctx.adapter ? await ctx.adapter.getExamples({
+      input: chunkInputStr.slice(0, 1e3),
+      taskType,
+      taskId
+    }) : [];
+    const key = fastHash(
+      stringify({
+        taskId,
+        taskType,
+        input: chunkInputStr,
+        instructions
+      })
+    );
+    const exactMatch = examples.find((x) => x.key === key);
+    if (exactMatch && exactMatch.output) {
+      return exactMatch.output;
+    }
+    const elementsText = chunk.map((elem, i) => {
+      const truncated = tokenizer.truncate(elem.stringified, options.tokensPerItem);
+      return `\u25A0${i}: ${truncated}\u25A0`;
+    }).join("\n");
+    const criteriaText = criteriaKeys.map((key2) => {
+      const criterion = sortingCriteria[key2];
+      const labelsText = criterion.labels.join(";");
+      return `**${key2}**: ${labelsText}`;
+    }).join("\n");
+    const exampleMessages = [];
+    for (const example of examples.slice(0, 3)) {
+      try {
+        const exampleInput = JSON.parse(example.input);
+        const exampleItems = Array.isArray(exampleInput) ? exampleInput : [exampleInput];
+        exampleMessages.push({
+          type: "text",
+          role: "user",
+          content: `Expert Example - Items to score:
+${exampleItems.map((el, i) => `\u25A0${i}: ${stringify(el, false).slice(0, 200)}\u25A0`).join("\n")}
+
+Score each item.`
+        });
+        const exampleOutput = example.output;
+        if (Array.isArray(exampleOutput) && exampleOutput.length > 0) {
+          const formattedScores = exampleOutput.map((score) => {
+            const pairs = criteriaKeys.map((key2) => `${key2}=${score.scores[key2] ?? 0}`).join(";");
+            return `\u25A0${score.elementIndex}:${pairs}\u25A0`;
+          }).join("\n");
+          exampleMessages.push({
+            type: "text",
+            role: "assistant",
+            content: `${formattedScores}
+${END}`
+          });
+          if (example.explanation) {
+            exampleMessages.push({
+              type: "text",
+              role: "assistant",
+              content: `Reasoning: ${example.explanation}`
+            });
+          }
+        }
+      } catch {
+      }
+    }
+    const { extracted } = await ctx.generateContent({
+      systemPrompt: `You are ranking items for sorting using ordered label arrays.
+
+${criteriaText}
+
+Instructions: "${instructions}"
+
+SCORING RULES:
+- For each item and each criterion, assign ONE label from the ordered list
+- Labels are ordered: first label = appears FIRST in sorted result, last label = appears LAST
+- Choose the label that best describes each item
+
+Output format:
+\u25A00:criterion1=label;criterion2=label\u25A0
+\u25A01:criterion1=label;criterion2=label\u25A0
+${END}
+
+IMPORTANT:
+- Rank every item (\u25A00 to \u25A0${chunkSize - 1})
+- Use exact criterion names: ${criteriaKeys.join(", ")}
+- Use exact labels from the lists above (lowercase, underscores)
+- Use semicolons (;) between criteria
+- Use equals (=) between criterion and label`,
+      stopSequences: [END],
+      messages: [
+        ...exampleMessages,
+        {
+          type: "text",
+          role: "user",
+          content: `Items to rank (\u25A00 to \u25A0${chunkSize - 1}):
+${elementsText}
+
+Rank each item using the labeled scales.
+Output format: \u25A0index:criterion1=label;criterion2=label\u25A0
+${END}`
+        }
+      ],
+      transform: (text) => {
+        const results = [];
+        const regex = /■(\d+):([^■]+)■/g;
+        let match;
+        while ((match = regex.exec(text)) !== null) {
+          const idx = parseInt(match[1] ?? "", 10);
+          const labelsStr = match[2] ?? "";
+          if (isNaN(idx) || idx < 0 || idx >= chunkSize) continue;
+          const scores = {};
+          let total = 0;
+          const pairs = labelsStr.split(";").filter((x) => x.trim().length > 0);
+          for (const pair of pairs) {
+            const [criterion, labelStr] = pair.split("=").map((x) => x.trim().toLowerCase().replace(/\s+/g, "_"));
+            if (!criterion || !labelStr) continue;
+            const labels = sortingCriteria[criterion]?.labels ?? [];
+            const labelIndex = labels.findIndex((l) => l === labelStr);
+            if (labelIndex >= 0) {
+              scores[criterion] = labelIndex;
+              total += labelIndex;
+            } else {
+              const middleIndex = labels.length > 0 ? Math.floor(labels.length / 2) : 5;
+              scores[criterion] = middleIndex;
+              total += middleIndex;
+            }
+          }
+          results[idx] = {
+            elementIndex: chunk[idx].index,
+            scores,
+            totalScore: total
+          };
+        }
+        for (let i = 0; i < chunkSize; i++) {
+          if (!results[i]) {
+            const scores = {};
+            let total = 0;
+            for (const key2 of criteriaKeys) {
+              const labels = sortingCriteria[key2]?.labels ?? [];
+              const middleIndex = labels.length > 0 ? Math.floor(labels.length / 2) : 5;
+              scores[key2] = middleIndex;
+              total += middleIndex;
+            }
+            results[i] = {
+              elementIndex: chunk[i].index,
+              scores,
+              totalScore: total
+            };
+          }
+        }
+        return results;
+      }
+    });
+    return extracted;
+  };
+  const limit = pLimit(10);
+  const chunkPromises = chunks.map((chunk) => limit(() => scoreChunk(chunk)));
+  const allScores = await Promise.all(chunkPromises);
+  const scoreMap = /* @__PURE__ */ new Map();
+  for (const chunkScores of allScores) {
+    for (const itemScore of chunkScores) {
+      const existing = scoreMap.get(itemScore.elementIndex);
+      if (existing) {
+        for (const key of criteriaKeys) {
+          existing.scores[key] = (existing.scores[key] + (itemScore.scores[key] ?? 0)) / 2;
+        }
+        existing.totalScore = (existing.totalScore + itemScore.totalScore) / 2;
+      } else {
+        scoreMap.set(itemScore.elementIndex, {
+          scores: { ...itemScore.scores },
+          totalScore: itemScore.totalScore
+        });
+      }
+    }
+  }
+  if (scoreMap.size !== input.length) {
+    throw new Error(`Score map size mismatch: expected ${input.length}, got ${scoreMap.size}`);
+  }
+  const scoreGroups = /* @__PURE__ */ new Map();
+  for (const [index, scoreData] of scoreMap.entries()) {
+    const roundedScore = Math.round(scoreData.totalScore * 100);
+    const group = scoreGroups.get(roundedScore) ?? [];
+    group.push(index);
+    scoreGroups.set(roundedScore, group);
+  }
+  const tiedGroups = Array.from(scoreGroups.values()).filter((group) => group.length > 1);
+  if (tiedGroups.length > 0) {
+    const tieBreakLimit = pLimit(10);
+    await Promise.all(
+      tiedGroups.map(
+        (tiedIndices) => tieBreakLimit(async () => {
+          if (tiedIndices.length <= 1) return;
+          const tiedElements = tiedIndices.map((idx) => elements[idx]);
+          const tieBreakText = tiedElements.map((elem, i) => {
+            const truncated = tokenizer.truncate(elem.stringified, options.tokensPerItem);
+            return `\u25A0${i}: ${truncated}\u25A0`;
+          }).join("\n");
+          const { extracted: tieBreakOrder } = await ctx.generateContent({
+            systemPrompt: `You are breaking a tie between items with identical total scores.
+
+Instructions: ${instructions}
+
+Criteria:
+${criteriaKeys.map((key) => {
+              const labels = sortingCriteria[key].labels.join(";");
+              return `- ${key}: ${labels}`;
+            }).join("\n")}
+
+Order these ${tiedElements.length} items from FIRST to LAST based on the instructions.
+Earlier labels in each criterion should come FIRST.
+
+Output format:
+\u25A0original_index\u25A0
+\u25A0original_index\u25A0
+${END}
+
+Output the indices in the order they should appear (first item at top).`,
+            stopSequences: [END],
+            messages: [
+              {
+                type: "text",
+                role: "user",
+                content: `Items with identical scores (need tie-breaking):
+${tieBreakText}
+
+Order them from first to last.
+Output format: \u25A0index\u25A0 (one per line)
+${END}`
+              }
+            ],
+            transform: (text) => {
+              const order = [];
+              const regex = /■(\d+)■/g;
+              let match;
+              while ((match = regex.exec(text)) !== null) {
+                const idx = parseInt(match[1] ?? "", 10);
+                if (!isNaN(idx) && idx >= 0 && idx < tiedElements.length) {
+                  order.push(idx);
+                }
+              }
+              for (let i = 0; i < tiedElements.length; i++) {
+                if (!order.includes(i)) {
+                  order.push(i);
+                }
+              }
+              return order;
+            }
+          });
+          for (let i = 0; i < tieBreakOrder.length; i++) {
+            const elementIndex = tiedElements[tieBreakOrder[i]].index;
+            const scoreData = scoreMap.get(elementIndex);
+            if (scoreData) {
+              scoreData.tieBreakOrder = i;
+            }
+          }
+        })
+      )
+    );
+  }
+  const sorted = Array.from(scoreMap.entries()).sort((a, b) => {
+    const scoreDiff = a[1].totalScore - b[1].totalScore;
+    if (scoreDiff !== 0) return scoreDiff;
+    const orderA = a[1].tieBreakOrder ?? 0;
+    const orderB = b[1].tieBreakOrder ?? 0;
+    return orderA - orderB;
+  }).map(([index]) => elements[index].element);
+  const result = sorted;
+  if (taskId && ctx.adapter && !ctx.controller.signal.aborted) {
+    const key = fastHash(
+      stringify({
+        taskId,
+        taskType,
+        input: JSON.stringify(input),
+        instructions
+      })
+    );
+    await ctx.adapter.saveExample({
+      key,
+      taskType,
+      taskId,
+      input: JSON.stringify(input),
+      output: result,
+      instructions,
+      metadata: {
+        cost: { input: 0, output: 0 },
+        latency: 0,
+        model: ctx.modelId,
+        tokens: { input: 0, output: 0 }
+      }
+    });
+  }
+  return result;
+};
+Zai.prototype.sort = function(input, instructions, _options) {
+  const context = new ZaiContext({
+    client: this.client,
+    modelId: this.Model,
+    taskId: this.taskId,
+    taskType: "zai.sort",
+    adapter: this.adapter
+  });
+  return new Response(
+    context,
+    sort(input, instructions, _options, context),
+    (result) => result
+    // Simplified form is just the sorted array
+  );
+};