@botpress/zai 1.0.1 → 1.2.0

package/dist/operations/label.js

@@ -0,0 +1,249 @@
+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();
+  await this.fetchModelDetails();
+  const taskId = this.taskId;
+  const taskType = "zai.label";
+  const TOTAL_MAX_TOKENS = clamp(options.chunkLength, 1e3, this.ModelDetails.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, meta } = 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: {
+        cost: {
+          input: meta.cost.input,
+          output: meta.cost.output
+        },
+        latency: meta.latency,
+        model: this.Model,
+        tokens: {
+          input: meta.tokens.input,
+          output: meta.tokens.output
+        }
+      },
+      input: inputAsString,
+      output: final
+    });
+  }
+  return convertToAnswer(final);
+};

package/dist/operations/rewrite.js

@@ -0,0 +1,123 @@
+import { z } from "@bpinternal/zui";
+import { fastHash, stringify, takeUntilTokens } from "../utils";
+import { Zai } from "../zai";
+import { PROMPT_INPUT_BUFFER } from "./constants";
+const Example = z.object({
+  input: z.string(),
+  output: z.string()
+});
+const Options = z.object({
+  examples: z.array(Example).default([]),
+  length: z.number().min(10).max(16e3).optional().describe("The maximum number of tokens to generate")
+});
+const START = "\u25A0START\u25A0";
+const END = "\u25A0END\u25A0";
+Zai.prototype.rewrite = async function(original, prompt, _options) {
+  const options = Options.parse(_options ?? {});
+  const tokenizer = await this.getTokenizer();
+  await this.fetchModelDetails();
+  const taskId = this.taskId;
+  const taskType = "zai.rewrite";
+  const INPUT_COMPONENT_SIZE = Math.max(100, (this.ModelDetails.input.maxTokens - PROMPT_INPUT_BUFFER) / 2);
+  prompt = tokenizer.truncate(prompt, INPUT_COMPONENT_SIZE);
+  const inputSize = tokenizer.count(original) + tokenizer.count(prompt);
+  const maxInputSize = this.ModelDetails.input.maxTokens - tokenizer.count(prompt) - PROMPT_INPUT_BUFFER;
+  if (inputSize > maxInputSize) {
+    throw new Error(
+      `The input size is ${inputSize} tokens long, which is more than the maximum of ${maxInputSize} tokens for this model (${this.ModelDetails.name} = ${this.ModelDetails.input.maxTokens} tokens)`
+    );
+  }
+  const instructions = [];
+  const originalSize = tokenizer.count(original);
+  if (options.length && originalSize > options.length) {
+    instructions.push(`The original text is ${originalSize} tokens long \u2013 it should be less than ${options.length}`);
+    instructions.push(
+      `The text must be standalone and complete in less than ${options.length} tokens, so it has to be shortened to fit the length as well`
+    );
+  }
+  const format = (before, prompt2) => {
+    return `
+Prompt: ${prompt2}
+
+${START}
+${before}
+${END}
+`.trim();
+  };
+  const Key = fastHash(
+    stringify({
+      taskId,
+      taskType,
+      input: original,
+      prompt
+    })
+  );
+  const formatExample = ({ input, output: output2, instructions: instructions2 }) => {
+    return [
+      { type: "text", role: "user", content: format(input, instructions2 || prompt) },
+      { type: "text", role: "assistant", content: `${START}${output2}${END}` }
+    ];
+  };
+  const defaultExamples = [
+    { input: "Hello, how are you?", output: "Bonjour, comment \xE7a va?", instructions: "translate to French" },
+    { input: "1\n2\n3", output: "3\n2\n1", instructions: "reverse the order" }
+  ];
+  const tableExamples = taskId ? await this.adapter.getExamples({
+    input: original,
+    taskId,
+    taskType
+  }) : [];
+  const exactMatch = tableExamples.find((x) => x.key === Key);
+  if (exactMatch) {
+    return exactMatch.output;
+  }
+  const savedExamples = [
+    ...tableExamples.map((x) => ({ input: x.input, output: x.output })),
+    ...options.examples
+  ];
+  const REMAINING_TOKENS = this.ModelDetails.input.maxTokens - tokenizer.count(prompt) - PROMPT_INPUT_BUFFER;
+  const examples = takeUntilTokens(
+    savedExamples.length ? savedExamples : defaultExamples,
+    REMAINING_TOKENS,
+    (el) => tokenizer.count(stringify(el.input)) + tokenizer.count(stringify(el.output))
+  ).map(formatExample).flat();
+  const { output, meta } = await this.callModel({
+    systemPrompt: `
+Rewrite the text between the ${START} and ${END} tags to match the user prompt.
+${instructions.map((x) => `\u2022 ${x}`).join("\n")}
+`.trim(),
+    messages: [...examples, { type: "text", content: format(original, prompt), role: "user" }],
+    maxTokens: options.length,
+    stopSequences: [END]
+  });
+  let result = output.choices[0]?.content;
+  if (result.includes(START)) {
+    result = result.slice(result.indexOf(START) + START.length);
+  }
+  if (result.includes(END)) {
+    result = result.slice(0, result.indexOf(END));
+  }
+  if (taskId) {
+    await this.adapter.saveExample({
+      key: Key,
+      metadata: {
+        cost: {
+          input: meta.cost.input,
+          output: meta.cost.output
+        },
+        latency: meta.latency,
+        model: this.Model,
+        tokens: {
+          input: meta.tokens.input,
+          output: meta.tokens.output
+        }
+      },
+      instructions: prompt,
+      input: original,
+      output: result,
+      taskType,
+      taskId
+    });
+  }
+  return result;
+};

package/dist/operations/summarize.js

@@ -0,0 +1,133 @@
+import { z } from "@bpinternal/zui";
+import { chunk } from "lodash-es";
+import { Zai } from "../zai";
+import { PROMPT_INPUT_BUFFER, PROMPT_OUTPUT_BUFFER } from "./constants";
+const Options = z.object({
+  prompt: z.string().describe("What should the text be summarized to?").default("New information, concepts and ideas that are deemed important"),
+  format: z.string().describe("How to format the example text").default(
+    "A normal text with multiple sentences and paragraphs. Use markdown to format the text into sections. Use headings, lists, and other markdown features to make the text more readable. Do not include links, images, or other non-text elements."
+  ),
+  length: z.number().min(10).max(1e5).describe("The length of the summary in tokens").default(250),
+  intermediateFactor: z.number().min(1).max(10).describe("How many times longer (than final length) are the intermediate summaries generated").default(4),
+  maxIterations: z.number().min(1).default(100),
+  sliding: z.object({
+    window: z.number().min(10).max(1e5),
+    overlap: z.number().min(0).max(1e5)
+  }).describe("Sliding window options").default({ window: 5e4, overlap: 250 })
+});
+const START = "\u25A0START\u25A0";
+const END = "\u25A0END\u25A0";
+Zai.prototype.summarize = async function(original, _options) {
+  const options = Options.parse(_options ?? {});
+  const tokenizer = await this.getTokenizer();
+  await this.fetchModelDetails();
+  const INPUT_COMPONENT_SIZE = Math.max(100, (this.ModelDetails.input.maxTokens - PROMPT_INPUT_BUFFER) / 4);
+  options.prompt = tokenizer.truncate(options.prompt, INPUT_COMPONENT_SIZE);
+  options.format = tokenizer.truncate(options.format, INPUT_COMPONENT_SIZE);
+  const maxOutputSize = this.ModelDetails.output.maxTokens - PROMPT_OUTPUT_BUFFER;
+  if (options.length > maxOutputSize) {
+    throw new Error(
+      `The desired output length is ${maxOutputSize} tokens long, which is more than the maximum of ${this.ModelDetails.output.maxTokens} tokens for this model (${this.ModelDetails.name})`
+    );
+  }
+  options.sliding.window = Math.min(options.sliding.window, this.ModelDetails.input.maxTokens - PROMPT_INPUT_BUFFER);
+  options.sliding.overlap = Math.min(options.sliding.overlap, options.sliding.window - 3 * options.sliding.overlap);
+  const format = (summary, newText) => {
+    return `
+${START}
+${summary.length ? summary : "<summary still empty>"}
+${END}
+
+Please amend the summary between the ${START} and ${END} tags to accurately reflect the prompt and the additional text below.
+
+<|start_new_information|>
+${newText}
+<|new_information|>`.trim();
+  };
+  const tokens = tokenizer.split(original);
+  const parts = Math.ceil(tokens.length / (options.sliding.window - options.sliding.overlap));
+  let iteration = 0;
+  const N = 2;
+  const useMergeSort = parts >= Math.pow(2, N);
+  const chunkSize = Math.ceil(tokens.length / (parts * N));
+  if (useMergeSort) {
+    const chunks = chunk(tokens, chunkSize).map((x) => x.join(""));
+    const allSummaries = await Promise.all(chunks.map((chunk2) => this.summarize(chunk2, options)));
+    return this.summarize(allSummaries.join("\n\n============\n\n"), options);
+  }
+  const summaries = [];
+  let currentSummary = "";
+  for (let i = 0; i < tokens.length; i += options.sliding.window) {
+    const from = Math.max(0, i - options.sliding.overlap);
+    const to = Math.min(tokens.length, i + options.sliding.window + options.sliding.overlap);
+    const isFirst = i === 0;
+    const isLast = to >= tokens.length;
+    const slice = tokens.slice(from, to).join("");
+    if (iteration++ >= options.maxIterations) {
+      break;
+    }
+    const instructions = [
+      `At each step, you will receive a part of the text to summarize. Make sure to reply with the new summary in the tags ${START} and ${END}.`,
+      "Summarize the text and make sure that the main points are included.",
+      "Ignore any unnecessary details and focus on the main points.",
+      "Use short and concise sentences to increase readability and information density.",
+      "When looking at the new information, focus on: " + options.prompt
+    ];
+    if (isFirst) {
+      instructions.push(
+        "The current summary is empty. You need to generate a summary that covers the main points of the text."
+      );
+    }
+    let generationLength = options.length;
+    if (!isLast) {
+      generationLength = Math.min(
+        tokenizer.count(currentSummary) + options.length * options.intermediateFactor,
+        maxOutputSize
+      );
+      instructions.push(
+        "You need to amend the summary to include the new information. Make sure the summary is complete and covers all the main points."
+      );
+      instructions.push(`The current summary is ${currentSummary.length} tokens long.`);
+      instructions.push(`You can amend the summary to be up to ${generationLength} tokens long.`);
+    }
+    if (isLast) {
+      instructions.push(
+        "This is the last part you will have to summarize. Make sure the summary is complete and covers all the main points."
+      );
+      instructions.push(
+        `The current summary is ${currentSummary.length} tokens long. You need to make sure it is ${options.length} tokens or less.`
+      );
+      if (currentSummary.length > options.length) {
+        instructions.push(
+          `The current summary is already too long, so you need to shorten it to ${options.length} tokens while also including the new information.`
+        );
+      }
+    }
+    const { output } = await this.callModel({
+      systemPrompt: `
+You are summarizing a text. The text is split into ${parts} parts, and you are currently working on part ${iteration}.
+At every step, you will receive the current summary and a new part of the text. You need to amend the summary to include the new information (if needed).
+The summary needs to cover the main points of the text and must be concise.
+
+IMPORTANT INSTRUCTIONS:
+${instructions.map((x) => `- ${x.trim()}`).join("\n")}
+
+FORMAT OF THE SUMMARY:
+${options.format}
+`.trim(),
+      messages: [{ type: "text", content: format(currentSummary, slice), role: "user" }],
+      maxTokens: generationLength,
+      stopSequences: [END]
+    });
+    let result = output?.choices[0]?.content;
+    if (result.includes(START)) {
+      result = result.slice(result.indexOf(START) + START.length);
+    }
+    if (result.includes("\u25A0")) {
+      result = result.slice(0, result.indexOf("\u25A0"));
+    }
+    summaries.push(result);
+    currentSummary = result;
+  }
+  return currentSummary.trim();
+};

package/dist/operations/text.js

@@ -0,0 +1,47 @@
+import { z } from "@bpinternal/zui";
+import { clamp } from "lodash-es";
+import { Zai } from "../zai";
+import { PROMPT_INPUT_BUFFER, PROMPT_OUTPUT_BUFFER } from "./constants";
+const Options = z.object({
+  length: z.number().min(1).max(1e5).optional().describe("The maximum number of tokens to generate")
+});
+Zai.prototype.text = async function(prompt, _options) {
+  const options = Options.parse(_options ?? {});
+  const tokenizer = await this.getTokenizer();
+  await this.fetchModelDetails();
+  prompt = tokenizer.truncate(prompt, Math.max(this.ModelDetails.input.maxTokens - PROMPT_INPUT_BUFFER, 100));
+  if (options.length) {
+    options.length = Math.min(this.ModelDetails.output.maxTokens - PROMPT_OUTPUT_BUFFER, options.length);
+  }
+  const instructions = [];
+  let chart = "";
+  if (options.length) {
+    const length = clamp(options.length * 0.75, 5, options.length);
+    instructions.push(`IMPORTANT: Length constraint: ${length} tokens/words`);
+    instructions.push(`The text must be standalone and complete in less than ${length} tokens/words`);
+  }
+  if (options.length && options.length <= 500) {
+    chart = `
+| Tokens      | Text Length (approximate)            |
+|-------------|--------------------------------------|
+| < 5 tokens  | 1-3 words                            |
+| 5-10 tokens | 3-6 words                            |
+| 10-20 tokens| 6-15 words                           |
+| 20-50 tokens| A short sentence (15-30 words)       |
+| 50-100 tokens| A medium sentence (30-70 words)     |
+| 100-200 tokens| A short paragraph (70-150 words)   |
+| 200-300 tokens| A medium paragraph (150-200 words) |
+| 300-500 tokens| A long paragraph (200-300 words)   |`.trim();
+  }
+  const { output } = await this.callModel({
+    systemPrompt: `
+Generate a text that fulfills the user prompt below. Answer directly to the prompt, without any acknowledgements or fluff. Also, make sure the text is standalone and complete.
+${instructions.map((x) => `- ${x}`).join("\n")}
+${chart}
+`.trim(),
+    temperature: 0.7,
+    messages: [{ type: "text", content: prompt, role: "user" }],
+    maxTokens: options.length
+  });
+  return output?.choices?.[0]?.content;
+};

package/dist/utils.js

@@ -0,0 +1,37 @@
+import { z } from "@bpinternal/zui";
+export const stringify = (input, beautify = true) => {
+  return typeof input === "string" && !!input.length ? input : input ? JSON.stringify(input, beautify ? null : void 0, beautify ? 2 : void 0) : "<input is null, false, undefined or empty>";
+};
+export function fastHash(str) {
+  let hash = 0;
+  for (let i = 0; i < str.length; i++) {
+    hash = (hash << 5) - hash + str.charCodeAt(i);
+    hash |= 0;
+  }
+  return (hash >>> 0).toString(16);
+}
+export const takeUntilTokens = (arr, tokens, count) => {
+  const result = [];
+  let total = 0;
+  for (const value of arr) {
+    const valueTokens = count(value);
+    if (total + valueTokens > tokens) {
+      break;
+    }
+    total += valueTokens;
+    result.push(value);
+  }
+  return result;
+};
+export const GenerationMetadata = z.object({
+  model: z.string(),
+  cost: z.object({
+    input: z.number(),
+    output: z.number()
+  }).describe("Cost in $USD"),
+  latency: z.number().describe("Latency in milliseconds"),
+  tokens: z.object({
+    input: z.number(),
+    output: z.number()
+  }).describe("Number of tokens used")
+});

package/dist/zai.js

@@ -0,0 +1,100 @@
+import { Cognitive } from "@botpress/cognitive";
+import { getWasmTokenizer } from "@bpinternal/thicktoken";
+import { z } from "@bpinternal/zui";
+import { TableAdapter } from "./adapters/botpress-table";
+import { MemoryAdapter } from "./adapters/memory";
+const ActiveLearning = z.object({
+  enable: z.boolean().describe("Whether to enable active learning").default(false),
+  tableName: z.string().regex(
+    /^[A-Za-z0-9_/-]{1,100}Table$/,
+    "Namespace must be alphanumeric and contain only letters, numbers, underscores, hyphens and slashes"
+  ).describe("The name of the table to store active learning tasks").default("ActiveLearningTable"),
+  taskId: z.string().regex(
+    /^[A-Za-z0-9_/-]{1,100}$/,
+    "Namespace must be alphanumeric and contain only letters, numbers, underscores, hyphens and slashes"
+  ).describe("The ID of the task").default("default")
+});
+const ZaiConfig = z.object({
+  client: z.custom(),
+  userId: z.string().describe("The ID of the user consuming the API").optional(),
+  modelId: z.custom(
+    (value) => {
+      if (typeof value !== "string") {
+        return false;
+      }
+      if (value !== "best" && value !== "fast" && !value.includes(":")) {
+        return false;
+      }
+      return true;
+    },
+    {
+      message: "Invalid model ID"
+    }
+  ).describe("The ID of the model you want to use").default("best"),
+  activeLearning: ActiveLearning.default({ enable: false }),
+  namespace: z.string().regex(
+    /^[A-Za-z0-9_/-]{1,100}$/,
+    "Namespace must be alphanumeric and contain only letters, numbers, underscores, hyphens and slashes"
+  ).default("zai")
+});
+export class Zai {
+  static tokenizer = null;
+  client;
+  _originalConfig;
+  _userId;
+  Model;
+  ModelDetails;
+  namespace;
+  adapter;
+  activeLearning;
+  constructor(config) {
+    this._originalConfig = config;
+    const parsed = ZaiConfig.parse(config);
+    this.client = Cognitive.isCognitiveClient(parsed.client) ? parsed.client : new Cognitive({ client: parsed.client });
+    this.namespace = parsed.namespace;
+    this._userId = parsed.userId;
+    this.Model = parsed.modelId;
+    this.activeLearning = parsed.activeLearning;
+    this.adapter = parsed.activeLearning?.enable ? new TableAdapter({ client: this.client.client, tableName: parsed.activeLearning.tableName }) : new MemoryAdapter([]);
+  }
+  /** @internal */
+  async callModel(props) {
+    return this.client.generateContent({
+      ...props,
+      model: this.Model,
+      userId: this._userId
+    });
+  }
+  async getTokenizer() {
+    Zai.tokenizer ??= await (async () => {
+      while (!getWasmTokenizer) {
+        await new Promise((resolve) => setTimeout(resolve, 25));
+      }
+      return getWasmTokenizer();
+    })();
+    return Zai.tokenizer;
+  }
+  async fetchModelDetails() {
+    if (!this.ModelDetails) {
+      this.ModelDetails = await this.client.getModelDetails(this.Model);
+    }
+  }
+  get taskId() {
+    if (!this.activeLearning.enable) {
+      return void 0;
+    }
+    return `${this.namespace}/${this.activeLearning.taskId}`.replace(/\/+/g, "/");
+  }
+  with(options) {
+    return new Zai({
+      ...this._originalConfig,
+      ...options
+    });
+  }
+  learn(taskId) {
+    return new Zai({
+      ...this._originalConfig,
+      activeLearning: { ...this.activeLearning, taskId, enable: true }
+    });
+  }
+}