@botpress/zai 2.0.14 → 2.0.16

package/dist/index.d.ts

@@ -172,16 +172,17 @@ type Options$1 = {
     instructions?: string;
     /** The maximum number of tokens per chunk */
     chunkLength?: number;
+    /** Whether to strictly follow the schema or not */
+    strict?: boolean;
 };
 type __Z<T extends any = any> = {
     _output: T;
 };
 type OfType<O, T extends __Z = __Z<O>> = T extends __Z<O> ? T : never;
-type AnyObjectOrArray = Record<string, unknown> | Array<unknown>;
 declare module '@botpress/zai' {
     interface Zai {
         /** Extracts one or many elements from an arbitrary input */
-        extract<S extends OfType<AnyObjectOrArray>>(input: unknown, schema: S, options?: Options$1): Promise<S['_output']>;
+        extract<S extends OfType<any>>(input: unknown, schema: S, options?: Options$1): Promise<S['_output']>;
     }
 }
 

package/dist/operations/extract.js

@@ -8,7 +8,8 @@ import { PROMPT_INPUT_BUFFER } from "./constants";
 import { JsonParsingError } from "./errors";
 const Options = z.object({
   instructions: z.string().optional().describe("Instructions to guide the user on how to extract the data"),
-  chunkLength: z.number().min(100).max(1e5).optional().describe("The maximum number of tokens per chunk").default(16e3)
+  chunkLength: z.number().min(100).max(1e5).optional().describe("The maximum number of tokens per chunk").default(16e3),
+  strict: z.boolean().optional().default(true).describe("Whether to strictly follow the schema or not")
 });
 const START = "\u25A0json_start\u25A0";
 const END = "\u25A0json_end\u25A0";
@@ -22,22 +23,34 @@ Zai.prototype.extract = async function(input, _schema, _options) {
   const taskType = "zai.extract";
   const PROMPT_COMPONENT = Math.max(this.ModelDetails.input.maxTokens - PROMPT_INPUT_BUFFER, 100);
   let isArrayOfObjects = false;
+  let wrappedValue = false;
   const originalSchema = schema;
   const baseType = (schema.naked ? schema.naked() : schema)?.constructor?.name ?? "unknown";
-  if (baseType === "ZodObject") {
-  } else if (baseType === "ZodArray") {
+  if (baseType === "ZodArray") {
+    isArrayOfObjects = true;
     let elementType = schema.element;
     if (elementType.naked) {
       elementType = elementType.naked();
     }
     if (elementType?.constructor?.name === "ZodObject") {
-      isArrayOfObjects = true;
       schema = elementType;
     } else {
-      throw new Error("Schema must be a ZodObject or a ZodArray<ZodObject>");
+      wrappedValue = true;
+      schema = z.object({
+        value: elementType
+      });
+    }
+  } else if (baseType !== "ZodObject") {
+    wrappedValue = true;
+    schema = z.object({
+      value: originalSchema
+    });
+  }
+  if (!options.strict) {
+    try {
+      schema = schema.partial();
+    } catch {
     }
-  } else {
-    throw new Error("Schema must be either a ZuiObject or a ZuiArray<ZuiObject>");
   }
   const schemaTypescript = schema.toTypescriptType({ declaration: false });
   const schemaLength = tokenizer.count(schemaTypescript);
@@ -46,16 +59,38 @@ Zai.prototype.extract = async function(input, _schema, _options) {
     this.ModelDetails.input.maxTokens - PROMPT_INPUT_BUFFER - schemaLength
   );
   const keys = Object.keys(schema.shape);
-  let inputAsString = stringify(input);
+  const inputAsString = stringify(input);
   if (tokenizer.count(inputAsString) > options.chunkLength) {
-    if (isArrayOfObjects) {
-      const tokens = tokenizer.split(inputAsString);
-      const chunks = chunk(tokens, options.chunkLength).map((x) => x.join(""));
-      const all = await Promise.all(chunks.map((chunk2) => this.extract(chunk2, originalSchema)));
-      return all.flat();
-    } else {
-      inputAsString = tokenizer.truncate(stringify(input), options.chunkLength);
-    }
+    const tokens = tokenizer.split(inputAsString);
+    const chunks = chunk(tokens, options.chunkLength).map((x) => x.join(""));
+    const all = await Promise.allSettled(
+      chunks.map(
+        (chunk2) => this.extract(chunk2, originalSchema, {
+          ...options,
+          strict: false
+          // We don't want to fail on strict mode for sub-chunks
+        })
+      )
+    ).then(
+      (results) => results.filter((x) => x.status === "fulfilled").map((x) => x.value)
+    );
+    const rows = all.map((x, idx) => `<part-${idx + 1}>
+${stringify(x, true)}
+</part-${idx + 1}>`).join("\n");
+    return this.extract(
+      `
+The result has been split into ${all.length} parts. Recursively merge the result into the final result.
+When merging arrays, take unique values.
+When merging conflictual (but defined) information, take the most reasonable and frequent value.
+Non-defined values are OK and normal. Don't delete fields because of null values. Focus on defined values.
+
+Here's the data:
+${rows}
+
+Merge it back into a final result.`.trim(),
+      originalSchema,
+      options
+    );
   }
   const instructions = [];
   if (options.instructions) {
@@ -76,6 +111,9 @@ Zai.prototype.extract = async function(input, _schema, _options) {
     instructions.push("You may have exactly one element in the input.");
     instructions.push(`The element must be a JSON object with exactly the format: ${START}${shape}${END}`);
   }
+  if (!options.strict) {
+    instructions.push("You may ignore any fields that are not present in the input. All keys are optional.");
+  }
   const EXAMPLES_TOKENS = PROMPT_COMPONENT - tokenizer.count(inputAsString) - tokenizer.count(instructions.join("\n"));
   const Key = fastHash(
     JSON.stringify({
@@ -188,13 +226,20 @@ ${instructions.map((x) => `\u2022 ${x}`).join("\n")}
       }
     ]
   });
-  const answer = output.choices[0]?.content;
-  const elements = answer.split(START).filter((x) => x.trim().length > 0).map((x) => {
+  const answer = output.choices[0]?.content ?? "{}";
+  const elements = answer?.split(START).filter((x) => x.trim().length > 0 && x.includes("}")).map((x) => {
     try {
       const json = x.slice(0, x.indexOf(END)).trim();
       const repairedJson = jsonrepair(json);
       const parsedJson = JSON5.parse(repairedJson);
-      return schema.parse(parsedJson);
+      const safe = schema.safeParse(parsedJson);
+      if (safe.success) {
+        return safe.data;
+      }
+      if (options.strict) {
+        throw new JsonParsingError(x, safe.error);
+      }
+      return parsedJson;
     } catch (error) {
       throw new JsonParsingError(x, error instanceof Error ? error : new Error("Unknown error"));
     }
@@ -203,10 +248,17 @@ ${instructions.map((x) => `\u2022 ${x}`).join("\n")}
   if (isArrayOfObjects) {
     final = elements;
   } else if (elements.length === 0) {
-    final = schema.parse({});
+    final = options.strict ? schema.parse({}) : {};
   } else {
     final = elements[0];
   }
+  if (wrappedValue) {
+    if (Array.isArray(final)) {
+      final = final.map((x) => "value" in x ? x.value : x);
+    } else {
+      final = "value" in final ? final.value : final;
+    }
+  }
   if (taskId) {
     await this.adapter.saveExample({
       key: Key,