@botpress/zai 1.0.1 → 1.1.0

package/dist/operations/extract.js

@@ -0,0 +1,212 @@
+import { z } from "@bpinternal/zui";
+import JSON5 from "json5";
+import { jsonrepair } from "jsonrepair";
+import { chunk, isArray } from "lodash-es";
+import { fastHash, stringify, takeUntilTokens } from "../utils";
+import { Zai } from "../zai";
+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)
+});
+const START = "\u25A0json_start\u25A0";
+const END = "\u25A0json_end\u25A0";
+const NO_MORE = "\u25A0NO_MORE_ELEMENT\u25A0";
+Zai.prototype.extract = async function(input, schema, _options) {
+  const options = Options.parse(_options ?? {});
+  const tokenizer = await this.getTokenizer();
+  const taskId = this.taskId;
+  const taskType = "zai.extract";
+  const PROMPT_COMPONENT = Math.max(this.Model.input.maxTokens - PROMPT_INPUT_BUFFER, 100);
+  let isArrayOfObjects = false;
+  const originalSchema = schema;
+  if (schema instanceof z.ZodObject) {
+  } else if (schema instanceof z.ZodArray) {
+    if (schema._def.type instanceof z.ZodObject) {
+      isArrayOfObjects = true;
+      schema = schema._def.type;
+    } else {
+      throw new Error("Schema must be a ZodObject or a ZodArray<ZodObject>");
+    }
+  } else {
+    throw new Error("Schema must be either a ZuiObject or a ZuiArray<ZuiObject>");
+  }
+  const schemaTypescript = schema.toTypescript({ declaration: false });
+  const schemaLength = tokenizer.count(schemaTypescript);
+  options.chunkLength = Math.min(options.chunkLength, this.Model.input.maxTokens - PROMPT_INPUT_BUFFER - schemaLength);
+  const keys = Object.keys(schema.shape);
+  let 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 instructions = [];
+  if (options.instructions) {
+    instructions.push(options.instructions);
+  }
+  const shape = `{ ${keys.map((key) => `"${key}": ...`).join(", ")} }`;
+  const abbv = "{ ... }";
+  if (isArrayOfObjects) {
+    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(`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 two elements, the output should look like this: ${START}${abbv}${END}${START}${abbv}${END}${NO_MORE}`
+    );
+  } else {
+    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}`);
+  }
+  const EXAMPLES_TOKENS = PROMPT_COMPONENT - tokenizer.count(inputAsString) - tokenizer.count(instructions.join("\n"));
+  const Key = fastHash(
+    JSON.stringify({
+      taskType,
+      taskId,
+      input: inputAsString,
+      instructions: options.instructions
+    })
+  );
+  const examples = taskId ? await this.adapter.getExamples({
+    input: inputAsString,
+    taskType,
+    taskId
+  }) : [];
+  const exactMatch = examples.find((x) => x.key === Key);
+  if (exactMatch) {
+    return exactMatch.output;
+  }
+  const defaultExample = isArrayOfObjects ? {
+    input: `The story goes as follow.
+Once upon a time, there was a person named Alice who was 30 years old.
+Then, there was a person named Bob who was 25 years old.
+The end.`,
+    schema: "Array<{ name: string, age: number }>",
+    instructions: "Extract all people",
+    extracted: [
+      {
+        name: "Alice",
+        age: 30
+      },
+      {
+        name: "Bob",
+        age: 25
+      }
+    ]
+  } : {
+    input: `The story goes as follow.
+Once upon a time, there was a person named Alice who was 30 years old.
+The end.`,
+    schema: "{ name: string, age: number }",
+    instructions: "Extract the person",
+    extracted: { name: "Alice", age: 30 }
+  };
+  const userExamples = examples.map((e) => ({
+    input: e.input,
+    extracted: e.output,
+    schema: schemaTypescript,
+    instructions: options.instructions
+  }));
+  let exampleId = 1;
+  const formatInput = (input2, schema2, instructions2) => {
+    const header = userExamples.length ? `Expert Example #${exampleId++}` : "Here's an example to help you understand the format:";
+    return `
+${header}
+
+<|start_schema|>
+${schema2}
+<|end_schema|>
+
+<|start_instructions|>
+${instructions2 ?? "No specific instructions, just follow the schema above."}
+<|end_instructions|>
+
+<|start_input|>
+${input2.trim()}
+<|end_input|>
+  `.trim();
+  };
+  const formatOutput = (extracted) => {
+    extracted = isArray(extracted) ? extracted : [extracted];
+    return extracted.map(
+      (x) => `
+${START}
+${JSON.stringify(x, null, 2)}
+${END}`.trim()
+    ).join("\n") + NO_MORE;
+  };
+  const formatExample = (example) => [
+    {
+      type: "text",
+      content: formatInput(stringify(example.input ?? null), example.schema, example.instructions),
+      role: "user"
+    },
+    {
+      type: "text",
+      content: formatOutput(example.extracted),
+      role: "assistant"
+    }
+  ];
+  const allExamples = takeUntilTokens(
+    userExamples.length ? userExamples : [defaultExample],
+    EXAMPLES_TOKENS,
+    (el) => tokenizer.count(stringify(el.input)) + tokenizer.count(stringify(el.extracted))
+  ).map(formatExample).flat();
+  const output = await this.callModel({
+    systemPrompt: `
+Extract the following information from the input:
+${schemaTypescript}
+====
+
+${instructions.map((x) => `\u2022 ${x}`).join("\n")}
+`.trim(),
+    stopSequences: [isArrayOfObjects ? NO_MORE : END],
+    messages: [
+      ...allExamples,
+      {
+        role: "user",
+        type: "text",
+        content: formatInput(inputAsString, schemaTypescript, options.instructions ?? "")
+      }
+    ]
+  });
+  const answer = output.choices[0]?.content;
+  const elements = answer.split(START).filter((x) => x.trim().length > 0).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);
+    } catch (error) {
+      throw new JsonParsingError(x, error instanceof Error ? error : new Error("Unknown error"));
+    }
+  }).filter((x) => x !== null);
+  let final;
+  if (isArrayOfObjects) {
+    final = elements;
+  } else if (elements.length === 0) {
+    final = schema.parse({});
+  } else {
+    final = elements[0];
+  }
+  if (taskId) {
+    await this.adapter.saveExample({
+      key: Key,
+      taskId: `zai/${taskId}`,
+      taskType,
+      instructions: options.instructions ?? "No specific instructions",
+      input: inputAsString,
+      output: final,
+      metadata: output.metadata
+    });
+  }
+  return final;
+};

package/dist/operations/filter.js

@@ -0,0 +1,179 @@
+import { z } from "@bpinternal/zui";
+import { clamp } from "lodash-es";
+import { fastHash, stringify, takeUntilTokens } from "../utils";
+import { Zai } from "../zai";
+import { PROMPT_INPUT_BUFFER, PROMPT_OUTPUT_BUFFER } from "./constants";
+const Example = z.object({
+  input: z.any(),
+  filter: z.boolean(),
+  reason: z.string().optional()
+});
+const Options = z.object({
+  tokensPerItem: z.number().min(1).max(1e5).optional().describe("The maximum number of tokens per item").default(250),
+  examples: z.array(Example).describe("Examples to filter the condition against").default([])
+});
+const END = "\u25A0END\u25A0";
+Zai.prototype.filter = async function(input, condition, _options) {
+  const options = Options.parse(_options ?? {});
+  const tokenizer = await this.getTokenizer();
+  const taskId = this.taskId;
+  const taskType = "zai.filter";
+  const MAX_ITEMS_PER_CHUNK = 50;
+  const TOKENS_TOTAL_MAX = this.Model.input.maxTokens - PROMPT_INPUT_BUFFER - PROMPT_OUTPUT_BUFFER;
+  const TOKENS_EXAMPLES_MAX = Math.floor(Math.max(250, TOKENS_TOTAL_MAX * 0.5));
+  const TOKENS_CONDITION_MAX = clamp(TOKENS_TOTAL_MAX * 0.25, 250, tokenizer.count(condition));
+  const TOKENS_INPUT_ARRAY_MAX = TOKENS_TOTAL_MAX - TOKENS_EXAMPLES_MAX - TOKENS_CONDITION_MAX;
+  condition = tokenizer.truncate(condition, TOKENS_CONDITION_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_INPUT_ARRAY_MAX || currentChunk.length >= MAX_ITEMS_PER_CHUNK) {
+      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 formatInput = (input2, condition2) => {
+    return `
+Condition to check:
+${condition2}
+
+Items (from \u25A00 to \u25A0${input2.length - 1})
+==============================
+${input2.map((x, idx) => `\u25A0${idx} = ${stringify(x.input ?? null, false)}`).join("\n")}
+`.trim();
+  };
+  const formatExamples = (examples) => {
+    return `
+${examples.map((x, idx) => `\u25A0${idx}:${!!x.filter ? "true" : "false"}`).join("")}
+${END}
+====
+Here's the reasoning behind each example:
+${examples.map((x, idx) => `\u25A0${idx}:${!!x.filter ? "true" : "false"}:${x.reason ?? "No reason provided"}`).join("\n")}
+`.trim();
+  };
+  const genericExamples = [
+    {
+      input: "apple",
+      filter: true,
+      reason: "Apples are fruits"
+    },
+    {
+      input: "Apple Inc.",
+      filter: false,
+      reason: "Apple Inc. is a company, not a fruit"
+    },
+    {
+      input: "banana",
+      filter: true,
+      reason: "Bananas are fruits"
+    },
+    {
+      input: "potato",
+      filter: false,
+      reason: "Potatoes are vegetables"
+    }
+  ];
+  const genericExamplesMessages = [
+    {
+      type: "text",
+      content: formatInput(genericExamples, "is a fruit"),
+      role: "user"
+    },
+    {
+      type: "text",
+      content: formatExamples(genericExamples),
+      role: "assistant"
+    }
+  ];
+  const filterChunk = async (chunk) => {
+    const examples = taskId ? await this.adapter.getExamples({
+      // The Table API can't search for a huge input string
+      input: JSON.stringify(chunk).slice(0, 1e3),
+      taskType,
+      taskId
+    }).then(
+      (x) => x.map((y) => ({ filter: y.output, input: y.input, reason: y.explanation }))
+    ) : [];
+    const allExamples = takeUntilTokens(
+      [...examples, ...options.examples ?? []],
+      TOKENS_EXAMPLES_MAX,
+      (el) => tokenizer.count(stringify(el.input))
+    );
+    const exampleMessages = [
+      {
+        type: "text",
+        content: formatInput(allExamples, condition),
+        role: "user"
+      },
+      {
+        type: "text",
+        content: formatExamples(allExamples),
+        role: "assistant"
+      }
+    ];
+    const output = await this.callModel({
+      systemPrompt: `
+You are given a list of items. Your task is to filter out the items that meet the condition below.
+You need to return the full list of items with the format:
+\u25A0x:true\u25A0y:false\u25A0z:true (where x, y, z are the indices of the items in the list)
+You need to start with "\u25A00" and go up to the last index "\u25A0${chunk.length - 1}".
+If an item meets the condition, you should return ":true", otherwise ":false".
+
+IMPORTANT: Make sure to read the condition and the examples carefully before making your decision.
+The condition is: "${condition}"
+`.trim(),
+      stopSequences: [END],
+      messages: [
+        ...exampleMessages.length ? exampleMessages : genericExamplesMessages,
+        {
+          type: "text",
+          content: formatInput(
+            chunk.map((x) => ({ input: x })),
+            condition
+          ),
+          role: "user"
+        }
+      ]
+    });
+    const answer = output.choices[0]?.content;
+    const indices = answer.trim().split("\u25A0").filter((x) => x.length > 0).map((x) => {
+      const [idx, filter] = x.split(":");
+      return { idx: parseInt(idx?.trim() ?? ""), filter: filter?.toLowerCase().trim() === "true" };
+    });
+    const partial = chunk.filter((_, idx) => {
+      return indices.find((x) => x.idx === idx)?.filter ?? false;
+    });
+    if (taskId) {
+      const key = fastHash(
+        stringify({
+          taskId,
+          taskType,
+          input: JSON.stringify(chunk),
+          condition
+        })
+      );
+      await this.adapter.saveExample({
+        key,
+        taskType,
+        taskId,
+        input: JSON.stringify(chunk),
+        output: partial,
+        instructions: condition,
+        metadata: output.metadata
+      });
+    }
+    return partial;
+  };
+  const filteredChunks = await Promise.all(chunks.map(filterChunk));
+  return filteredChunks.flat();
+};

package/dist/operations/label.js

@@ -0,0 +1,237 @@
+import { z } from "@bpinternal/zui";
+import { clamp, chunk } from "lodash-es";
+import { fastHash, stringify, takeUntilTokens } from "../utils";
+import { Zai } from "../zai";
+import { PROMPT_INPUT_BUFFER } from "./constants";
+const LABELS = {
+  ABSOLUTELY_NOT: "ABSOLUTELY_NOT",
+  PROBABLY_NOT: "PROBABLY_NOT",
+  AMBIGUOUS: "AMBIGUOUS",
+  PROBABLY_YES: "PROBABLY_YES",
+  ABSOLUTELY_YES: "ABSOLUTELY_YES"
+};
+const ALL_LABELS = Object.values(LABELS).join(" | ");
+const Options = z.object({
+  examples: z.array(
+    z.object({
+      input: z.any(),
+      labels: z.record(z.object({ label: z.enum(ALL_LABELS), explanation: z.string().optional() }))
+    })
+  ).default([]).describe("Examples to help the user make a decision"),
+  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)
+});
+const Labels = z.record(z.string().min(1).max(250), z.string()).superRefine((labels, ctx) => {
+  const keys = Object.keys(labels);
+  for (const key of keys) {
+    if (key.length < 1 || key.length > 250) {
+      ctx.addIssue({ message: `The label key "${key}" must be between 1 and 250 characters long`, code: "custom" });
+    }
+    if (keys.lastIndexOf(key) !== keys.indexOf(key)) {
+      ctx.addIssue({ message: `Duplicate label: ${labels[key]}`, code: "custom" });
+    }
+    if (/[^a-zA-Z0-9_]/.test(key)) {
+      ctx.addIssue({
+        message: `The label key "${key}" must only contain alphanumeric characters and underscores`,
+        code: "custom"
+      });
+    }
+  }
+  return true;
+});
+const parseLabel = (label) => {
+  label = label.toUpperCase().replace(/\s+/g, "_").replace(/_{2,}/g, "_").trim();
+  if (label.includes("ABSOLUTELY") && label.includes("NOT")) {
+    return LABELS.ABSOLUTELY_NOT;
+  } else if (label.includes("NOT")) {
+    return LABELS.PROBABLY_NOT;
+  } else if (label.includes("AMBIGUOUS")) {
+    return LABELS.AMBIGUOUS;
+  }
+  if (label.includes("YES")) {
+    return LABELS.PROBABLY_YES;
+  } else if (label.includes("ABSOLUTELY") && label.includes("YES")) {
+    return LABELS.ABSOLUTELY_YES;
+  }
+  return LABELS.AMBIGUOUS;
+};
+Zai.prototype.label = async function(input, _labels, _options) {
+  const options = Options.parse(_options ?? {});
+  const labels = Labels.parse(_labels);
+  const tokenizer = await this.getTokenizer();
+  const taskId = this.taskId;
+  const taskType = "zai.label";
+  const TOTAL_MAX_TOKENS = clamp(options.chunkLength, 1e3, this.Model.input.maxTokens - PROMPT_INPUT_BUFFER);
+  const CHUNK_EXAMPLES_MAX_TOKENS = clamp(Math.floor(TOTAL_MAX_TOKENS * 0.5), 250, 1e4);
+  const CHUNK_INPUT_MAX_TOKENS = clamp(
+    TOTAL_MAX_TOKENS - CHUNK_EXAMPLES_MAX_TOKENS,
+    TOTAL_MAX_TOKENS * 0.5,
+    TOTAL_MAX_TOKENS
+  );
+  const inputAsString = stringify(input);
+  if (tokenizer.count(inputAsString) > CHUNK_INPUT_MAX_TOKENS) {
+    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) => this.label(chunk2, _labels)));
+    return allLabels.reduce((acc, x) => {
+      Object.keys(x).forEach((key) => {
+        if (acc[key] === true) {
+          acc[key] = true;
+        } else {
+          acc[key] = acc[key] || x[key];
+        }
+      });
+      return acc;
+    }, {});
+  }
+  const END = "\u25A0END\u25A0";
+  const Key = fastHash(
+    JSON.stringify({
+      taskType,
+      taskId,
+      input: inputAsString,
+      instructions: options.instructions ?? ""
+    })
+  );
+  const convertToAnswer = (mapping) => {
+    return Object.keys(labels).reduce((acc, key) => {
+      acc[key] = mapping[key]?.label === "ABSOLUTELY_YES" || mapping[key]?.label === "PROBABLY_YES";
+      return acc;
+    }, {});
+  };
+  const examples = taskId ? await this.adapter.getExamples({
+    input: inputAsString,
+    taskType,
+    taskId
+  }) : [];
+  options.examples.forEach((example) => {
+    examples.push({
+      key: fastHash(JSON.stringify(example)),
+      input: example.input,
+      similarity: 1,
+      explanation: "",
+      output: example.labels
+    });
+  });
+  const exactMatch = examples.find((x) => x.key === Key);
+  if (exactMatch) {
+    return convertToAnswer(exactMatch.output);
+  }
+  const allExamples = takeUntilTokens(
+    examples,
+    CHUNK_EXAMPLES_MAX_TOKENS,
+    (el) => tokenizer.count(stringify(el.input)) + tokenizer.count(stringify(el.output)) + tokenizer.count(el.explanation ?? "") + 100
+  ).map((example, idx) => [
+    {
+      type: "text",
+      role: "user",
+      content: `
+Expert Example #${idx + 1}
+
+<|start_input|>
+${stringify(example.input)}
+<|end_input|>`.trim()
+    },
+    {
+      type: "text",
+      role: "assistant",
+      content: `
+Expert Example #${idx + 1}
+============
+${Object.keys(example.output).map(
+        (key) => `
+\u25A0${key}:\u3010${example.output[key]?.explanation}\u3011:${example.output[key]?.label}\u25A0
+`.trim()
+      ).join("\n")}
+${END}
+`.trim()
+    }
+  ]).flat();
+  const format = Object.keys(labels).map((key) => {
+    return `
+\u25A0${key}:\u3010explanation (where "explanation" is answering the question "${labels[key]}")\u3011:x\u25A0 (where x is ${ALL_LABELS})
+`.trim();
+  }).join("\n\n");
+  const output = await this.callModel({
+    stopSequences: [END],
+    systemPrompt: `
+You need to tag the input with the following labels based on the question asked:
+${LABELS.ABSOLUTELY_NOT}: You are absolutely sure that the answer is "NO" to the question.
+${LABELS.PROBABLY_NOT}: You are leaning towards "NO" to the question.
+${LABELS.AMBIGUOUS}: You are unsure about the answer to the question.
+${LABELS.PROBABLY_YES}: You are leaning towards "YES" to the question.
+${LABELS.ABSOLUTELY_YES}: You are absolutely sure that the answer is "YES" to the question.
+
+You need to return a mapping of the labels, an explanation and the answer for each label following the format below:
+\`\`\`
+${format}
+${END}
+\`\`\`
+
+${options.instructions}
+
+===
+You should consider the Expert Examples below to help you make your decision.
+In your "Analysis", please refer to the Expert Examples # to justify your decision.
+`.trim(),
+    messages: [
+      ...allExamples,
+      {
+        type: "text",
+        role: "user",
+        content: `
+Input to tag:
+<|start_input|>
+${inputAsString}
+<|end_input|>
+
+Answer with this following format:
+\`\`\`
+${format}
+${END}
+\`\`\`
+
+Format cheatsheet:
+\`\`\`
+\u25A0label:\u3010explanation\u3011:x\u25A0
+\`\`\`
+
+Where \`x\` is one of the following: ${ALL_LABELS}
+
+Remember: In your \`explanation\`, please refer to the Expert Examples # (and quote them) that are relevant to ground your decision-making process.
+The Expert Examples are there to help you make your decision. They have been provided by experts in the field and their answers (and reasoning) are considered the ground truth and should be used as a reference to make your decision when applicable.
+For example, you can say: "According to Expert Example #1, ..."`.trim()
+      }
+    ]
+  });
+  const answer = output.choices[0].content;
+  const final = Object.keys(labels).reduce((acc, key) => {
+    const match = answer.match(new RegExp(`\u25A0${key}:\u3010(.+)\u3011:(\\w{2,})\u25A0`, "i"));
+    if (match) {
+      const explanation = match[1].trim();
+      const label = parseLabel(match[2]);
+      acc[key] = {
+        explanation,
+        label
+      };
+    } else {
+      acc[key] = {
+        explanation: "",
+        label: LABELS.AMBIGUOUS
+      };
+    }
+    return acc;
+  }, {});
+  if (taskId) {
+    await this.adapter.saveExample({
+      key: Key,
+      taskType,
+      taskId,
+      instructions: options.instructions ?? "",
+      metadata: output.metadata,
+      input: inputAsString,
+      output: final
+    });
+  }
+  return convertToAnswer(final);
+};