@botpress/zai 2.2.0 → 2.3.0

package/dist/operations/rate.js

@@ -0,0 +1,350 @@
+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 RATING_VALUES = {
+  very_bad: 1,
+  bad: 2,
+  average: 3,
+  good: 4,
+  very_good: 5
+};
+const _Options = z.object({
+  tokensPerItem: z.number().min(1).max(1e5).optional().describe("The maximum number of tokens per item").default(250),
+  maxItemsPerChunk: z.number().min(1).max(100).optional().describe("The maximum number of items to rate per chunk").default(50)
+});
+const END = "\u25A0END\u25A0";
+const rate = async (input, instructions, _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 taskId = ctx.taskId;
+  const taskType = "zai.rate";
+  const TOKENS_TOTAL_MAX = model.input.maxTokens - PROMPT_INPUT_BUFFER - PROMPT_OUTPUT_BUFFER;
+  const isStringInstructions = typeof instructions === "string";
+  const criteriaKeys = isStringInstructions ? [] : Object.keys(instructions);
+  const generateCriteriaPrompt = isStringInstructions ? `Generate 3-5 evaluation criteria for: "${instructions}"
+
+For each criterion, provide 5 labels (very_bad, bad, average, good, very_good) with brief descriptions.
+
+Output format (JSON):
+{
+  "criterion1_name": {
+    "very_bad": "description",
+    "bad": "description",
+    "average": "description",
+    "good": "description",
+    "very_good": "description"
+  },
+  "criterion2_name": { ... }
+}
+
+Keep criterion names short (1-2 words, lowercase, use underscores).
+Keep descriptions brief (5-10 words each).` : `For these evaluation criteria, provide 5 labels (very_bad, bad, average, good, very_good) with brief descriptions for each:
+
+${criteriaKeys.map((key) => `- ${key}: ${instructions[key]}`).join("\n")}
+
+Output format (JSON):
+{
+  "${criteriaKeys[0]}": {
+    "very_bad": "description",
+    "bad": "description",
+    "average": "description",
+    "good": "description",
+    "very_good": "description"
+  }
+  ${criteriaKeys.length > 1 ? "..." : ""}
+}
+
+Keep descriptions brief (5-10 words each).`;
+  const { extracted: evaluationCriteria } = await ctx.generateContent({
+    systemPrompt: `You are creating evaluation criteria for rating items on a 1-5 scale.
+Each criterion must have exactly 5 labels: very_bad (1), bad (2), average (3), good (4), very_good (5).
+Output valid JSON only.`,
+    messages: [
+      {
+        type: "text",
+        role: "user",
+        content: generateCriteriaPrompt
+      }
+    ],
+    transform: (text) => {
+      const jsonMatch = text.match(/```(?:json)?\s*(\{[\s\S]*?\})\s*```/) || text.match(/(\{[\s\S]*\})/);
+      if (!jsonMatch) {
+        throw new Error("Failed to parse evaluation criteria JSON");
+      }
+      return JSON.parse(jsonMatch[1]);
+    }
+  });
+  const finalCriteriaKeys = Object.keys(evaluationCriteria);
+  if (finalCriteriaKeys.length === 0) {
+    throw new Error("No evaluation criteria generated");
+  }
+  const TOKENS_CRITERIA_MAX = Math.floor(TOKENS_TOTAL_MAX * 0.3);
+  const TOKENS_ITEMS_MAX = TOKENS_TOTAL_MAX - TOKENS_CRITERIA_MAX;
+  let chunks = [];
+  let currentChunk = [];
+  let currentChunkTokens = 0;
+  for (const element of input) {
+    const elementAsString = tokenizer.truncate(stringify(element, false), options.tokensPerItem);
+    const elementTokens = tokenizer.count(elementAsString);
+    if (currentChunkTokens + elementTokens > TOKENS_ITEMS_MAX || currentChunk.length >= options.maxItemsPerChunk) {
+      if (currentChunk.length > 0) {
+        chunks.push(currentChunk);
+      }
+      currentChunk = [];
+      currentChunkTokens = 0;
+    }
+    currentChunk.push(element);
+    currentChunkTokens += elementTokens;
+  }
+  if (currentChunk.length > 0) {
+    chunks.push(currentChunk);
+  }
+  chunks = chunks.filter((x) => x.length > 0);
+  const rateChunk = async (chunk) => {
+    ctx.controller.signal.throwIfAborted();
+    const chunkInputStr = JSON.stringify(chunk);
+    const examples = taskId && ctx.adapter ? await ctx.adapter.getExamples({
+      input: chunkInputStr.slice(0, 1e3),
+      // Limit search string length
+      taskType,
+      taskId
+    }) : [];
+    const key = fastHash(
+      stringify({
+        taskId,
+        taskType,
+        input: chunkInputStr,
+        instructions: stringify(instructions)
+      })
+    );
+    const exactMatch = examples.find((x) => x.key === key);
+    if (exactMatch && exactMatch.output) {
+      return {
+        ratings: exactMatch.output,
+        meta: { cost: { input: 0, output: 0 }, latency: 0, tokens: { input: 0, output: 0 } }
+      };
+    }
+    const formatCriteria = () => {
+      return finalCriteriaKeys.map((key2) => {
+        const labels = evaluationCriteria[key2];
+        return `**${key2}**:
+  - very_bad (1): ${labels?.very_bad}
+  - bad (2): ${labels?.bad}
+  - average (3): ${labels?.average}
+  - good (4): ${labels?.good}
+  - very_good (5): ${labels?.very_good}`;
+      }).join("\n\n");
+    };
+    const formatItems = (items) => {
+      return items.map((item, idx) => {
+        const itemStr = tokenizer.truncate(stringify(item, false), options.tokensPerItem);
+        return `\u25A0${idx}: ${itemStr}\u25A0`;
+      }).join("\n");
+    };
+    const exampleMessages = [];
+    for (const example of examples.slice(0, 5)) {
+      try {
+        const exampleInput = JSON.parse(example.input);
+        exampleMessages.push({
+          type: "text",
+          role: "user",
+          content: `Expert Example - Items to rate:
+${formatItems(Array.isArray(exampleInput) ? exampleInput : [exampleInput])}
+
+Rate each item on all criteria.`
+        });
+        const exampleOutput = example.output;
+        if (Array.isArray(exampleOutput) && exampleOutput.length > 0) {
+          const formattedRatings = exampleOutput.map((rating, idx) => {
+            const pairs = finalCriteriaKeys.map((key2) => {
+              const value = rating[key2];
+              if (typeof value === "number") {
+                const labelMap = {
+                  1: "very_bad",
+                  2: "bad",
+                  3: "average",
+                  4: "good",
+                  5: "very_good"
+                };
+                return `${key2}=${labelMap[value] || "average"}`;
+              }
+              return null;
+            }).filter(Boolean).join(";");
+            return `\u25A0${idx}:${pairs}\u25A0`;
+          }).join("\n");
+          exampleMessages.push({
+            type: "text",
+            role: "assistant",
+            content: `${formattedRatings}
+${END}`
+          });
+          if (example.explanation) {
+            exampleMessages.push({
+              type: "text",
+              role: "assistant",
+              content: `Reasoning: ${example.explanation}`
+            });
+          }
+        }
+      } catch {
+      }
+    }
+    const { extracted, meta } = await ctx.generateContent({
+      systemPrompt: `You are rating items based on evaluation criteria.
+
+Evaluation Criteria:
+${formatCriteria()}
+
+For each item, rate it on EACH criterion using one of these labels:
+very_bad, bad, average, good, very_good
+
+Output format:
+\u25A00:criterion1=label;criterion2=label;criterion3=label\u25A0
+\u25A01:criterion1=label;criterion2=label;criterion3=label\u25A0
+${END}
+
+IMPORTANT:
+- Rate every item (\u25A00 to \u25A0${chunk.length - 1})
+- Use exact criterion names: ${finalCriteriaKeys.join(", ")}
+- Use exact label names: very_bad, bad, average, good, very_good
+- Use semicolons (;) between criteria
+- Use equals (=) between criterion and label`,
+      stopSequences: [END],
+      messages: [
+        ...exampleMessages,
+        {
+          type: "text",
+          role: "user",
+          content: `Items to rate (\u25A00 to \u25A0${chunk.length - 1}):
+${formatItems(chunk)}
+
+Rate each item on all criteria.
+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 ratingsStr = match[2] ?? "";
+          if (isNaN(idx) || idx < 0 || idx >= chunk.length) {
+            continue;
+          }
+          const itemRatings = {};
+          let total = 0;
+          const pairs = ratingsStr.split(";").filter((x) => x.trim().length > 0);
+          for (const pair of pairs) {
+            const [criterion, label] = pair.split("=").map((x) => x.trim());
+            if (!criterion || !label) continue;
+            const labelLower = label.toLowerCase().replace(/\s+/g, "_");
+            const ratingValue = RATING_VALUES[labelLower] ?? 3;
+            itemRatings[criterion] = ratingValue;
+            total += ratingValue;
+          }
+          itemRatings.total = total;
+          results[idx] = itemRatings;
+        }
+        for (let i = 0; i < chunk.length; i++) {
+          if (!results[i]) {
+            const defaultRatings = {};
+            let total = 0;
+            for (const key2 of finalCriteriaKeys) {
+              defaultRatings[key2] = 3;
+              total += 3;
+            }
+            defaultRatings.total = total;
+            results[i] = defaultRatings;
+          }
+        }
+        return results;
+      }
+    });
+    return { ratings: extracted, meta };
+  };
+  const limit = pLimit(10);
+  const chunkPromises = chunks.map((chunk) => limit(() => rateChunk(chunk)));
+  const ratedChunks = await Promise.all(chunkPromises);
+  const allRatings = ratedChunks.flatMap((result) => result.ratings);
+  const totalMeta = ratedChunks.reduce(
+    (acc, result) => ({
+      cost: {
+        input: acc.cost.input + result.meta.cost.input,
+        output: acc.cost.output + result.meta.cost.output
+      },
+      latency: Math.max(acc.latency, result.meta.latency),
+      // Use max latency
+      tokens: {
+        input: acc.tokens.input + result.meta.tokens.input,
+        output: acc.tokens.output + result.meta.tokens.output
+      }
+    }),
+    {
+      cost: { input: 0, output: 0 },
+      latency: 0,
+      tokens: { input: 0, output: 0 }
+    }
+  );
+  if (taskId && ctx.adapter && !ctx.controller.signal.aborted) {
+    const key = fastHash(
+      stringify({
+        taskId,
+        taskType,
+        input: JSON.stringify(input),
+        instructions: stringify(instructions)
+      })
+    );
+    await ctx.adapter.saveExample({
+      key,
+      taskType,
+      taskId,
+      input: JSON.stringify(input),
+      output: allRatings,
+      instructions: typeof instructions === "string" ? instructions : JSON.stringify(instructions),
+      metadata: {
+        cost: {
+          input: totalMeta.cost.input,
+          output: totalMeta.cost.output
+        },
+        latency: totalMeta.latency,
+        model: ctx.modelId,
+        tokens: {
+          input: totalMeta.tokens.input,
+          output: totalMeta.tokens.output
+        }
+      }
+    });
+  }
+  return allRatings;
+};
+Zai.prototype.rate = function(input, instructions, _options) {
+  const context = new ZaiContext({
+    client: this.client,
+    modelId: this.Model,
+    taskId: this.taskId,
+    taskType: "zai.rate",
+    adapter: this.adapter
+  });
+  return new Response(
+    context,
+    rate(input, instructions, _options, context),
+    (result) => {
+      if (typeof instructions === "string") {
+        return result.map((r) => r.total);
+      }
+      return result;
+    }
+  );
+};