@botpress/zai 2.1.18 → 2.1.20

package/dist/index.d.ts

@@ -1,4 +1,4 @@
-import { Cognitive, Model, BotpressClientLike, GenerateContentInput, GenerateContentOutput } from '@botpress/cognitive';
+import { Cognitive, Models, Model, BotpressClientLike, GenerateContentInput, GenerateContentOutput } from '@botpress/cognitive';
 import { TextTokenizer } from '@bpinternal/thicktoken';
 
 type GenerationMetadata = {
@@ -41,7 +41,6 @@ declare abstract class Adapter {
     abstract saveExample<TInput, TOutput>(props: SaveExampleProps<TInput, TOutput>): Promise<void>;
 }
 
-type ModelId = Required<Parameters<Cognitive['generateContent']>[0]['model']>;
 type ActiveLearning = {
     enable: boolean;
     tableName: string;
@@ -50,7 +49,7 @@ type ActiveLearning = {
 type ZaiConfig = {
     client: BotpressClientLike | Cognitive;
     userId?: string;
-    modelId?: ModelId | string;
+    modelId?: Models;
     activeLearning?: ActiveLearning;
     namespace?: string;
 };
@@ -59,7 +58,7 @@ declare class Zai {
     protected client: Cognitive;
     private _originalConfig;
     private _userId;
-    protected Model: ModelId;
+    protected Model: Models;
     protected ModelDetails: Model;
     protected namespace: string;
     protected adapter: Adapter;

package/dist/operations/errors.js

@@ -1,15 +1,119 @@
+import { ZodError } from "@bpinternal/zui";
 export class JsonParsingError extends Error {
   constructor(json, error) {
-    const message = `Error parsing JSON:
-
----JSON---
-${json}
-
----Error---
-
- ${error}`;
+    const message = JsonParsingError._formatError(json, error);
     super(message);
     this.json = json;
     this.error = error;
   }
+  static _formatError(json, error) {
+    let errorMessage = "Error parsing JSON:\n\n";
+    errorMessage += `---JSON---
+${json}
+
+`;
+    if (error instanceof ZodError) {
+      errorMessage += "---Validation Errors---\n\n";
+      errorMessage += JsonParsingError._formatZodError(error);
+    } else {
+      errorMessage += "---Error---\n\n";
+      errorMessage += "The JSON provided is not valid JSON.\n";
+      errorMessage += `Details: ${error.message}
+`;
+    }
+    return errorMessage;
+  }
+  static _formatZodError(zodError) {
+    const issues = zodError.issues;
+    if (issues.length === 0) {
+      return "Unknown validation error\n";
+    }
+    let message = "";
+    for (let i = 0; i < issues.length; i++) {
+      const issue = issues[i];
+      const path = issue.path.length > 0 ? issue.path.join(".") : "root";
+      message += `${i + 1}. Field: "${path}"
+`;
+      switch (issue.code) {
+        case "invalid_type":
+          message += `   Problem: Expected ${issue.expected}, but received ${issue.received}
+`;
+          message += `   Message: ${issue.message}
+`;
+          break;
+        case "invalid_string":
+          if ("validation" in issue) {
+            message += `   Problem: Invalid ${issue.validation} format
+`;
+          }
+          message += `   Message: ${issue.message}
+`;
+          break;
+        case "too_small":
+          if (issue.type === "string") {
+            if (issue.exact) {
+              message += `   Problem: String must be exactly ${issue.minimum} characters
+`;
+            } else {
+              message += `   Problem: String must be at least ${issue.minimum} characters
+`;
+            }
+          } else if (issue.type === "number") {
+            message += `   Problem: Number must be ${issue.inclusive ? "at least" : "greater than"} ${issue.minimum}
+`;
+          } else if (issue.type === "array") {
+            message += `   Problem: Array must contain ${issue.inclusive ? "at least" : "more than"} ${issue.minimum} items
+`;
+          }
+          message += `   Message: ${issue.message}
+`;
+          break;
+        case "too_big":
+          if (issue.type === "string") {
+            if (issue.exact) {
+              message += `   Problem: String must be exactly ${issue.maximum} characters
+`;
+            } else {
+              message += `   Problem: String must be at most ${issue.maximum} characters
+`;
+            }
+          } else if (issue.type === "number") {
+            message += `   Problem: Number must be ${issue.inclusive ? "at most" : "less than"} ${issue.maximum}
+`;
+          } else if (issue.type === "array") {
+            message += `   Problem: Array must contain ${issue.inclusive ? "at most" : "fewer than"} ${issue.maximum} items
+`;
+          }
+          message += `   Message: ${issue.message}
+`;
+          break;
+        case "invalid_enum_value":
+          message += `   Problem: Invalid value "${issue.received}"
+`;
+          message += `   Allowed values: ${issue.options.map((o) => `"${o}"`).join(", ")}
+`;
+          message += `   Message: ${issue.message}
+`;
+          break;
+        case "invalid_literal":
+          message += `   Problem: Expected the literal value "${issue.expected}", but received "${issue.received}"
+`;
+          message += `   Message: ${issue.message}
+`;
+          break;
+        case "invalid_union":
+          message += "   Problem: Value doesn't match any of the expected formats\n";
+          message += `   Message: ${issue.message}
+`;
+          break;
+        default:
+          message += `   Problem: ${issue.message}
+`;
+      }
+      if (i < issues.length - 1) {
+        message += "\n";
+      }
+    }
+    return message;
+  }
 }

package/dist/operations/extract.js

@@ -1,7 +1,8 @@
-import { z } from "@bpinternal/zui";
+import { z, transforms } from "@bpinternal/zui";
 import JSON5 from "json5";
 import { jsonrepair } from "jsonrepair";
 import { chunk, isArray } from "lodash-es";
+import pLimit from "p-limit";
 import { ZaiContext } from "../context";
 import { Response } from "../response";
 import { getTokenizer } from "../tokenizer";
@@ -17,9 +18,10 @@ const Options = z.object({
 const START = "\u25A0json_start\u25A0";
 const END = "\u25A0json_end\u25A0";
 const NO_MORE = "\u25A0NO_MORE_ELEMENT\u25A0";
+const ZERO_ELEMENTS = "\u25A0ZERO_ELEMENTS\u25A0";
 const extract = async (input, _schema, _options, ctx) => {
   ctx.controller.signal.throwIfAborted();
-  let schema = _schema;
+  let schema = transforms.fromJSONSchema(transforms.toJSONSchema(_schema));
   const options = Options.parse(_options ?? {});
   const tokenizer = await getTokenizer();
   const model = await ctx.getModel();
@@ -62,19 +64,22 @@ const extract = async (input, _schema, _options, ctx) => {
   const keys = Object.keys(schema.shape);
   const inputAsString = stringify(input);
   if (tokenizer.count(inputAsString) > options.chunkLength) {
+    const limit = pLimit(10);
     const tokens = tokenizer.split(inputAsString);
     const chunks = chunk(tokens, options.chunkLength).map((x) => x.join(""));
     const all = await Promise.allSettled(
       chunks.map(
-        (chunk2) => extract(
-          chunk2,
-          originalSchema,
-          {
-            ...options,
-            strict: false
-            // We don't want to fail on strict mode for sub-chunks
-          },
-          ctx
+        (chunk2) => limit(
+          () => extract(
+            chunk2,
+            originalSchema,
+            {
+              ...options,
+              strict: false
+              // We don't want to fail on strict mode for sub-chunks
+            },
+            ctx
+          )
         )
       )
     ).then(
@@ -110,8 +115,11 @@ Merge it back into a final result.`.trim(),
     instructions.push("You may have multiple elements, or zero elements in the input.");
     instructions.push("You must extract each element separately.");
     instructions.push(`Each element must be a JSON object with exactly the format: ${START}${shape}${END}`);
+    instructions.push(`If there are no elements to extract, respond with ${ZERO_ELEMENTS}.`);
     instructions.push(`When you are done extracting all elements, type "${NO_MORE}" to finish.`);
-    instructions.push(`For example, if you have zero elements, the output should look like this: ${NO_MORE}`);
+    instructions.push(
+      `For example, if you have zero elements, the output should look like this: ${ZERO_ELEMENTS}${NO_MORE}`
+    );
     instructions.push(
       `For example, if you have two elements, the output should look like this: ${START}${abbv}${END}${START}${abbv}${END}${NO_MORE}`
     );

package/dist/operations/filter.js

@@ -1,5 +1,6 @@
 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";
@@ -191,7 +192,8 @@ The condition is: "${condition}"
     }
     return partial;
   };
-  const filteredChunks = await Promise.all(chunks.map(filterChunk));
+  const limit = pLimit(10);
+  const filteredChunks = await Promise.all(chunks.map((chunk) => limit(() => filterChunk(chunk))));
   return filteredChunks.flat();
 };
 Zai.prototype.filter = function(input, condition, _options) {

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 = [];