@botpress/zai 2.0.15 → 2.1.0

package/dist/operations/extract.js

@@ -2,56 +2,103 @@ import { z } from "@bpinternal/zui";
 import JSON5 from "json5";
 import { jsonrepair } from "jsonrepair";
 import { chunk, isArray } from "lodash-es";
+import { ZaiContext } from "../context";
+import { Response } from "../response";
+import { getTokenizer } from "../tokenizer";
 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)
+  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";
 const NO_MORE = "\u25A0NO_MORE_ELEMENT\u25A0";
-Zai.prototype.extract = async function(input, _schema, _options) {
+const extract = async (input, _schema, _options, ctx) => {
+  ctx.controller.signal.throwIfAborted();
   let schema = _schema;
   const options = Options.parse(_options ?? {});
-  const tokenizer = await this.getTokenizer();
-  await this.fetchModelDetails();
-  const taskId = this.taskId;
+  const tokenizer = await getTokenizer();
+  const model = await ctx.getModel();
+  const taskId = ctx.taskId;
   const taskType = "zai.extract";
-  const PROMPT_COMPONENT = Math.max(this.ModelDetails.input.maxTokens - PROMPT_INPUT_BUFFER, 100);
+  const PROMPT_COMPONENT = Math.max(model.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);
-  options.chunkLength = Math.min(
-    options.chunkLength,
-    this.ModelDetails.input.maxTokens - PROMPT_INPUT_BUFFER - schemaLength
-  );
+  options.chunkLength = Math.min(options.chunkLength, model.input.maxTokens - PROMPT_INPUT_BUFFER - schemaLength);
   const keys = Object.keys(schema.shape);
   const inputAsString = stringify(input);
   if (tokenizer.count(inputAsString) > options.chunkLength) {
     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 this.extract(all, originalSchema, options);
+    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
+        )
+      )
+    ).then(
+      (results) => results.filter((x) => x.status === "fulfilled").map((x) => x.value)
+    );
+    ctx.controller.signal.throwIfAborted();
+    const rows = all.map((x, idx) => `<part-${idx + 1}>
+${stringify(x, true)}
+</part-${idx + 1}>`).join("\n");
+    return 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,
+      ctx
+    );
   }
   const instructions = [];
   if (options.instructions) {
@@ -72,6 +119,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({
@@ -81,7 +131,7 @@ Zai.prototype.extract = async function(input, _schema, _options) {
       instructions: options.instructions
     })
   );
-  const examples = taskId ? await this.adapter.getExamples({
+  const examples = taskId && ctx.adapter ? await ctx.adapter.getExamples({
     input: inputAsString,
     taskType,
     taskId
@@ -140,9 +190,9 @@ ${input2.trim()}
 <|end_input|>
   `.trim();
   };
-  const formatOutput = (extracted) => {
-    extracted = isArray(extracted) ? extracted : [extracted];
-    return extracted.map(
+  const formatOutput = (extracted2) => {
+    extracted2 = isArray(extracted2) ? extracted2 : [extracted2];
+    return extracted2.map(
       (x) => `
 ${START}
 ${JSON.stringify(x, null, 2)}
@@ -166,7 +216,7 @@ ${END}`.trim()
     EXAMPLES_TOKENS,
     (el) => tokenizer.count(stringify(el.input)) + tokenizer.count(stringify(el.extracted))
   ).map(formatExample).flat();
-  const { output, meta } = await this.callModel({
+  const { meta, extracted } = await ctx.generateContent({
     systemPrompt: `
 Extract the following information from the input:
 ${schemaTypescript}
@@ -182,29 +232,42 @@ ${instructions.map((x) => `\u2022 ${x}`).join("\n")}
         type: "text",
         content: formatInput(inputAsString, schemaTypescript, options.instructions ?? "")
       }
-    ]
+    ],
+    transform: (text) => (text || "{}")?.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);
+        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"));
+      }
+    }).filter((x) => x !== null)
   });
-  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({});
+    final = extracted;
+  } else if (extracted.length === 0) {
+    final = options.strict ? schema.parse({}) : {};
   } else {
-    final = elements[0];
+    final = extracted[0];
   }
-  if (taskId) {
-    await this.adapter.saveExample({
+  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 && ctx.adapter && !ctx.controller.signal.aborted) {
+    await ctx.adapter.saveExample({
       key: Key,
       taskId: `zai/${taskId}`,
       taskType,
@@ -217,7 +280,7 @@ ${instructions.map((x) => `\u2022 ${x}`).join("\n")}
           output: meta.cost.output
         },
         latency: meta.latency,
-        model: this.Model,
+        model: ctx.modelId,
         tokens: {
           input: meta.tokens.input,
           output: meta.tokens.output
@@ -227,3 +290,13 @@ ${instructions.map((x) => `\u2022 ${x}`).join("\n")}
   }
   return final;
 };
+Zai.prototype.extract = function(input, schema, _options) {
+  const context = new ZaiContext({
+    client: this.client,
+    modelId: this.Model,
+    taskId: this.taskId,
+    taskType: "zai.extract",
+    adapter: this.adapter
+  });
+  return new Response(context, extract(input, schema, _options, context), (result) => result);
+};

package/dist/operations/filter.js

@@ -1,5 +1,8 @@
 import { z } from "@bpinternal/zui";
 import { clamp } from "lodash-es";
+import { ZaiContext } from "../context";
+import { Response } from "../response";
+import { getTokenizer } from "../tokenizer";
 import { fastHash, stringify, takeUntilTokens } from "../utils";
 import { Zai } from "../zai";
 import { PROMPT_INPUT_BUFFER, PROMPT_OUTPUT_BUFFER } from "./constants";
@@ -13,14 +16,15 @@ const _Options = z.object({
   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 filter = async (input, condition, _options, ctx) => {
+  ctx.controller.signal.throwIfAborted();
   const options = _Options.parse(_options ?? {});
-  const tokenizer = await this.getTokenizer();
-  await this.fetchModelDetails();
-  const taskId = this.taskId;
+  const tokenizer = await getTokenizer();
+  const model = await ctx.getModel();
+  const taskId = ctx.taskId;
   const taskType = "zai.filter";
   const MAX_ITEMS_PER_CHUNK = 50;
-  const TOKENS_TOTAL_MAX = this.ModelDetails.input.maxTokens - PROMPT_INPUT_BUFFER - PROMPT_OUTPUT_BUFFER;
+  const TOKENS_TOTAL_MAX = 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;
@@ -97,7 +101,7 @@ ${examples.map((x, idx) => `\u25A0${idx}:${!!x.filter ? "true" : "false"}:${x.re
     }
   ];
   const filterChunk = async (chunk) => {
-    const examples = taskId ? await this.adapter.getExamples({
+    const examples = taskId && ctx.adapter ? await ctx.adapter.getExamples({
       // The Table API can't search for a huge input string
       input: JSON.stringify(chunk).slice(0, 1e3),
       taskType,
@@ -122,7 +126,7 @@ ${examples.map((x, idx) => `\u25A0${idx}:${!!x.filter ? "true" : "false"}:${x.re
         role: "assistant"
       }
     ];
-    const { output, meta } = await this.callModel({
+    const { extracted: partial, meta } = await ctx.generateContent({
       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:
@@ -144,17 +148,18 @@ The condition is: "${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;
+      ],
+      transform: (text) => {
+        const indices = text.trim().split("\u25A0").filter((x) => x.length > 0).map((x) => {
+          const [idx, filter2] = x.split(":");
+          return { idx: parseInt(idx?.trim() ?? ""), filter: filter2?.toLowerCase().trim() === "true" };
+        });
+        return chunk.filter((_, idx) => {
+          return indices.find((x) => x.idx === idx && x.filter) ?? false;
+        });
+      }
     });
-    if (taskId) {
+    if (taskId && ctx.adapter && !ctx.controller.signal.aborted) {
       const key = fastHash(
         stringify({
           taskId,
@@ -163,7 +168,7 @@ The condition is: "${condition}"
           condition
         })
       );
-      await this.adapter.saveExample({
+      await ctx.adapter.saveExample({
         key,
         taskType,
         taskId,
@@ -176,7 +181,7 @@ The condition is: "${condition}"
             output: meta.cost.output
           },
           latency: meta.latency,
-          model: this.Model,
+          model: ctx.modelId,
           tokens: {
             input: meta.tokens.input,
             output: meta.tokens.output
@@ -189,3 +194,13 @@ The condition is: "${condition}"
   const filteredChunks = await Promise.all(chunks.map(filterChunk));
   return filteredChunks.flat();
 };
+Zai.prototype.filter = function(input, condition, _options) {
+  const context = new ZaiContext({
+    client: this.client,
+    modelId: this.Model,
+    taskId: this.taskId,
+    taskType: "zai.filter",
+    adapter: this.adapter
+  });
+  return new Response(context, filter(input, condition, _options, context), (result) => result);
+};

package/dist/operations/label.js

@@ -1,5 +1,8 @@
 import { z } from "@bpinternal/zui";
-import { clamp, chunk } from "lodash-es";
+import { chunk, clamp } from "lodash-es";
+import { ZaiContext } from "../context";
+import { Response } from "../response";
+import { getTokenizer } from "../tokenizer";
 import { fastHash, stringify, takeUntilTokens } from "../utils";
 import { Zai } from "../zai";
 import { PROMPT_INPUT_BUFFER } from "./constants";
@@ -39,24 +42,24 @@ const _Labels = z.record(z.string().min(1).max(250), z.string()).superRefine((la
   }
   return true;
 });
-const parseLabel = (label) => {
-  label = label.toUpperCase().replace(/\s+/g, "_").replace(/_{2,}/g, "_").trim();
-  if (label.includes("ABSOLUTELY") && label.includes("NOT")) {
+const parseLabel = (label2) => {
+  label2 = label2.toUpperCase().replace(/\s+/g, "_").replace(/_{2,}/g, "_").trim();
+  if (label2.includes("ABSOLUTELY") && label2.includes("NOT")) {
     return LABELS.ABSOLUTELY_NOT;
-  } else if (label.includes("NOT")) {
+  } else if (label2.includes("NOT")) {
     return LABELS.PROBABLY_NOT;
-  } else if (label.includes("AMBIGUOUS")) {
+  } else if (label2.includes("AMBIGUOUS")) {
     return LABELS.AMBIGUOUS;
   }
-  if (label.includes("YES")) {
+  if (label2.includes("YES")) {
     return LABELS.PROBABLY_YES;
-  } else if (label.includes("ABSOLUTELY") && label.includes("YES")) {
+  } else if (label2.includes("ABSOLUTELY") && label2.includes("YES")) {
     return LABELS.ABSOLUTELY_YES;
   }
   return LABELS.AMBIGUOUS;
 };
-const getConfidence = (label) => {
-  switch (label) {
+const getConfidence = (label2) => {
+  switch (label2) {
     case LABELS.ABSOLUTELY_NOT:
     case LABELS.ABSOLUTELY_YES:
       return 1;
@@ -67,14 +70,15 @@ const getConfidence = (label) => {
       return 0;
   }
 };
-Zai.prototype.label = async function(input, _labels, _options) {
+const label = async (input, _labels, _options, ctx) => {
+  ctx.controller.signal.throwIfAborted();
   const options = _Options.parse(_options ?? {});
   const labels = _Labels.parse(_labels);
-  const tokenizer = await this.getTokenizer();
-  await this.fetchModelDetails();
-  const taskId = this.taskId;
+  const tokenizer = await getTokenizer();
+  const model = await ctx.getModel();
+  const taskId = ctx.taskId;
   const taskType = "zai.label";
-  const TOTAL_MAX_TOKENS = clamp(options.chunkLength, 1e3, this.ModelDetails.input.maxTokens - PROMPT_INPUT_BUFFER);
+  const TOTAL_MAX_TOKENS = clamp(options.chunkLength, 1e3, 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,
@@ -85,7 +89,7 @@ Zai.prototype.label = async function(input, _labels, _options) {
   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)));
+    const allLabels = await Promise.all(chunks.map((chunk2) => label(chunk2, _labels, _options, ctx)));
     return allLabels.reduce((acc, x) => {
       Object.keys(x).forEach((key) => {
         if (acc[key]?.value === true) {
@@ -118,7 +122,7 @@ Zai.prototype.label = async function(input, _labels, _options) {
       return acc;
     }, {});
   };
-  const examples = taskId ? await this.adapter.getExamples({
+  const examples = taskId && ctx.adapter ? await ctx.adapter.getExamples({
     input: inputAsString,
     taskType,
     taskId
@@ -171,7 +175,7 @@ ${END}
 \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({
+  const { extracted, meta } = await ctx.generateContent({
     stopSequences: [END],
     systemPrompt: `
 You need to tag the input with the following labels based on the question asked:
@@ -221,28 +225,27 @@ Remember: In your \`explanation\`, please refer to the Expert Examples # (and qu
 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()
       }
-    ]
+    ],
+    transform: (text) => Object.keys(labels).reduce((acc, key) => {
+      const match = text.match(new RegExp(`\u25A0${key}:\u3010(.+)\u3011:(\\w{2,})\u25A0`, "i"));
+      if (match) {
+        const explanation = match[1].trim();
+        const label2 = parseLabel(match[2]);
+        acc[key] = {
+          explanation,
+          label: label2
+        };
+      } else {
+        acc[key] = {
+          explanation: "",
+          label: LABELS.AMBIGUOUS
+        };
+      }
+      return acc;
+    }, {})
   });
-  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({
+  if (taskId && ctx.adapter && !ctx.controller.signal.aborted) {
+    await ctx.adapter.saveExample({
       key: Key,
       taskType,
       taskId,
@@ -253,15 +256,35 @@ For example, you can say: "According to Expert Example #1, ..."`.trim()
           output: meta.cost.output
         },
         latency: meta.latency,
-        model: this.Model,
+        model: ctx.modelId,
         tokens: {
           input: meta.tokens.input,
           output: meta.tokens.output
         }
       },
       input: inputAsString,
-      output: final
+      output: extracted
     });
   }
-  return convertToAnswer(final);
+  return convertToAnswer(extracted);
+};
+Zai.prototype.label = function(input, labels, _options) {
+  const context = new ZaiContext({
+    client: this.client,
+    modelId: this.Model,
+    taskId: this.taskId,
+    taskType: "zai.label",
+    adapter: this.adapter
+  });
+  return new Response(
+    context,
+    label(input, labels, _options, context),
+    (result) => Object.keys(result).reduce(
+      (acc, key) => {
+        acc[key] = result[key].value;
+        return acc;
+      },
+      {}
+    )
+  );
 };

package/dist/operations/rewrite.js

@@ -1,4 +1,7 @@
 import { z } from "@bpinternal/zui";
+import { ZaiContext } from "../context";
+import { Response } from "../response";
+import { getTokenizer } from "../tokenizer";
 import { fastHash, stringify, takeUntilTokens } from "../utils";
 import { Zai } from "../zai";
 import { PROMPT_INPUT_BUFFER } from "./constants";
@@ -12,19 +15,20 @@ const Options = z.object({
 });
 const START = "\u25A0START\u25A0";
 const END = "\u25A0END\u25A0";
-Zai.prototype.rewrite = async function(original, prompt, _options) {
+const rewrite = async (original, prompt, _options, ctx) => {
+  ctx.controller.signal.throwIfAborted();
   const options = Options.parse(_options ?? {});
-  const tokenizer = await this.getTokenizer();
-  await this.fetchModelDetails();
-  const taskId = this.taskId;
+  const tokenizer = await getTokenizer();
+  const model = await ctx.getModel();
+  const taskId = ctx.taskId;
   const taskType = "zai.rewrite";
-  const INPUT_COMPONENT_SIZE = Math.max(100, (this.ModelDetails.input.maxTokens - PROMPT_INPUT_BUFFER) / 2);
+  const INPUT_COMPONENT_SIZE = Math.max(100, (model.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;
+  const maxInputSize = model.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)`
+      `The input size is ${inputSize} tokens long, which is more than the maximum of ${maxInputSize} tokens for this model (${model.name} = ${model.input.maxTokens} tokens)`
     );
   }
   const instructions = [];
@@ -52,17 +56,17 @@ ${END}
       prompt
     })
   );
-  const formatExample = ({ input, output: output2, instructions: instructions2 }) => {
+  const formatExample = ({ input, output, instructions: instructions2 }) => {
     return [
       { type: "text", role: "user", content: format(input, instructions2 || prompt) },
-      { type: "text", role: "assistant", content: `${START}${output2}${END}` }
+      { type: "text", role: "assistant", content: `${START}${output}${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({
+  const tableExamples = taskId && ctx.adapter ? await ctx.adapter.getExamples({
     input: original,
     taskId,
     taskType
@@ -75,30 +79,36 @@ ${END}
     ...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 REMAINING_TOKENS = model.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({
+  const { extracted, meta } = await ctx.generateContent({
     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]
+    stopSequences: [END],
+    transform: (text) => {
+      if (!text.trim().length) {
+        throw new Error("The model did not return a valid rewrite. The response was empty.");
+      }
+      return text;
+    }
   });
-  let result = output.choices[0]?.content;
+  let result = extracted;
   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({
+  if (taskId && ctx.adapter && !ctx.controller.signal.aborted) {
+    await ctx.adapter.saveExample({
       key: Key,
       metadata: {
         cost: {
@@ -106,7 +116,7 @@ ${instructions.map((x) => `\u2022 ${x}`).join("\n")}
           output: meta.cost.output
         },
         latency: meta.latency,
-        model: this.Model,
+        model: ctx.modelId,
         tokens: {
           input: meta.tokens.input,
           output: meta.tokens.output
@@ -121,3 +131,13 @@ ${instructions.map((x) => `\u2022 ${x}`).join("\n")}
   }
   return result;
 };
+Zai.prototype.rewrite = function(original, prompt, _options) {
+  const context = new ZaiContext({
+    client: this.client,
+    modelId: this.Model,
+    taskId: this.taskId,
+    taskType: "zai.rewrite",
+    adapter: this.adapter
+  });
+  return new Response(context, rewrite(original, prompt, _options, context), (result) => result);
+};