@botpress/zai 2.0.16 → 2.1.0

package/dist/context.js

@@ -0,0 +1,131 @@
+import { EventEmitter } from "./emitter";
+export class ZaiContext {
+  _startedAt = Date.now();
+  _inputCost = 0;
+  _outputCost = 0;
+  _inputTokens = 0;
+  _outputTokens = 0;
+  _totalCachedResponses = 0;
+  _totalRequests = 0;
+  _totalErrors = 0;
+  _totalResponses = 0;
+  taskId;
+  taskType;
+  modelId;
+  adapter;
+  source;
+  _eventEmitter;
+  controller = new AbortController();
+  _client;
+  constructor(props) {
+    this._client = props.client.clone();
+    this.taskId = props.taskId;
+    this.modelId = props.modelId;
+    this.adapter = props.adapter;
+    this.source = props.source;
+    this.taskType = props.taskType;
+    this._eventEmitter = new EventEmitter();
+    this._client.on("request", () => {
+      this._totalRequests++;
+      this._eventEmitter.emit("update", this.usage);
+    });
+    this._client.on("response", (_req, res) => {
+      this._totalResponses++;
+      if (res.meta.cached) {
+        this._totalCachedResponses++;
+      } else {
+        this._inputTokens += res.meta.tokens.input || 0;
+        this._outputTokens += res.meta.tokens.output || 0;
+        this._inputCost += res.meta.cost.input || 0;
+        this._outputCost += res.meta.cost.output || 0;
+      }
+      this._eventEmitter.emit("update", this.usage);
+    });
+    this._client.on("error", () => {
+      this._totalErrors++;
+      this._eventEmitter.emit("update", this.usage);
+    });
+  }
+  async getModel() {
+    return this._client.getModelDetails(this.modelId);
+  }
+  on(type, listener) {
+    this._eventEmitter.on(type, listener);
+    return this;
+  }
+  clear() {
+    this._eventEmitter.clear();
+  }
+  async generateContent(props) {
+    const maxRetries = Math.max(props.maxRetries ?? 3, 0);
+    const transform = props.transform;
+    let lastError = null;
+    const messages = [...props.messages || []];
+    for (let attempt = 0; attempt <= maxRetries; attempt++) {
+      try {
+        const response = await this._client.generateContent({
+          ...props,
+          messages,
+          signal: this.controller.signal,
+          model: this.modelId,
+          meta: {
+            integrationName: props.meta?.integrationName || "zai",
+            promptCategory: props.meta?.promptCategory || `zai:${this.taskType}`,
+            promptSource: props.meta?.promptSource || `zai:${this.taskType}:${this.taskId ?? "default"}`
+          }
+        });
+        const content = response.output.choices[0]?.content;
+        const str = typeof content === "string" ? content : content?.[0]?.text || "";
+        let output;
+        messages.push({
+          role: "assistant",
+          content: str || "<Invalid output, no content provided>"
+        });
+        if (!transform) {
+          output = str;
+        } else {
+          output = transform(str, response.output);
+        }
+        return { meta: response.meta, output: response.output, text: str, extracted: output };
+      } catch (error) {
+        lastError = error;
+        if (attempt === maxRetries) {
+          throw lastError;
+        }
+        messages.push({
+          role: "user",
+          content: `ERROR PARSING OUTPUT
+
+${lastError.message}.
+
+Please return a valid response addressing the error above.`
+        });
+      }
+    }
+    throw lastError;
+  }
+  get elapsedTime() {
+    return Date.now() - this._startedAt;
+  }
+  get usage() {
+    return {
+      requests: {
+        errors: this._totalErrors,
+        requests: this._totalRequests,
+        responses: this._totalResponses,
+        cached: this._totalCachedResponses,
+        percentage: this._totalRequests > 0 ? (this._totalResponses + this._totalErrors) / this._totalRequests : 0
+      },
+      tokens: {
+        input: this._inputTokens,
+        output: this._outputTokens,
+        total: this._inputTokens + this._outputTokens
+      },
+      cost: {
+        input: this._inputCost,
+        output: this._outputCost,
+        total: this._inputCost + this._outputCost
+      }
+    };
+  }
+}

package/dist/emitter.js

@@ -0,0 +1,42 @@
+export class EventEmitter {
+  _listeners = {};
+  emit(type, event) {
+    const listeners = this._listeners[type];
+    if (!listeners) {
+      return;
+    }
+    for (const listener of listeners) {
+      listener(event);
+    }
+  }
+  once(type, listener) {
+    const wrapped = (event) => {
+      this.off(type, wrapped);
+      listener(event);
+    };
+    this.on(type, wrapped);
+  }
+  on(type, listener) {
+    if (!this._listeners[type]) {
+      this._listeners[type] = [];
+    }
+    this._listeners[type].push(listener);
+  }
+  off(type, listener) {
+    const listeners = this._listeners[type];
+    if (!listeners) {
+      return;
+    }
+    const index = listeners.indexOf(listener);
+    if (index !== -1) {
+      listeners.splice(index, 1);
+    }
+  }
+  clear(type) {
+    if (type) {
+      delete this._listeners[type];
+    } else {
+      this._listeners = {};
+    }
+  }
+}

package/dist/index.d.ts

@@ -1,4 +1,4 @@
-import { Cognitive, Model, BotpressClientLike } from '@botpress/cognitive';
+import { Cognitive, Model, BotpressClientLike, GenerateContentInput, GenerateContentOutput } from '@botpress/cognitive';
 import { TextTokenizer } from '@bpinternal/thicktoken';
 
 type GenerationMetadata = {
@@ -74,6 +74,99 @@ declare class Zai {
     learn(taskId: string): Zai;
 }
 
+type Meta = Awaited<ReturnType<Cognitive['generateContent']>>['meta'];
+type GenerateContentProps<T> = Omit<GenerateContentInput, 'model' | 'signal'> & {
+    maxRetries?: number;
+    transform?: (text: string | undefined, output: GenerateContentOutput) => T;
+};
+type ZaiContextProps = {
+    client: Cognitive;
+    taskType: string;
+    taskId: string;
+    modelId: string;
+    adapter?: Adapter;
+    source?: GenerateContentInput['meta'];
+};
+type Usage = {
+    requests: {
+        requests: number;
+        errors: number;
+        responses: number;
+        cached: number;
+        percentage: number;
+    };
+    cost: {
+        input: number;
+        output: number;
+        total: number;
+    };
+    tokens: {
+        input: number;
+        output: number;
+        total: number;
+    };
+};
+type ContextEvents = {
+    update: Usage;
+};
+declare class ZaiContext {
+    private _startedAt;
+    private _inputCost;
+    private _outputCost;
+    private _inputTokens;
+    private _outputTokens;
+    private _totalCachedResponses;
+    private _totalRequests;
+    private _totalErrors;
+    private _totalResponses;
+    taskId: string;
+    taskType: string;
+    modelId: GenerateContentInput['model'];
+    adapter?: Adapter;
+    source?: GenerateContentInput['meta'];
+    private _eventEmitter;
+    controller: AbortController;
+    private _client;
+    constructor(props: ZaiContextProps);
+    getModel(): Promise<Model>;
+    on<K extends keyof ContextEvents>(type: K, listener: (event: ContextEvents[K]) => void): this;
+    clear(): void;
+    generateContent<Out = string>(props: GenerateContentProps<Out>): Promise<{
+        meta: Meta;
+        output: GenerateContentOutput;
+        text: string | undefined;
+        extracted: Out;
+    }>;
+    get elapsedTime(): number;
+    get usage(): Usage;
+}
+
+type ResponseEvents<TComplete = any> = {
+    progress: Usage;
+    complete: TComplete;
+    error: unknown;
+};
+declare class Response<T = any, S = T> implements PromiseLike<S> {
+    private _promise;
+    private _eventEmitter;
+    private _context;
+    private _elasped;
+    private _simplify;
+    constructor(context: ZaiContext, promise: Promise<T>, simplify: (value: T) => S);
+    on<K extends keyof ResponseEvents<T>>(type: K, listener: (event: ResponseEvents<T>[K]) => void): this;
+    off<K extends keyof ResponseEvents<T>>(type: K, listener: (event: ResponseEvents<T>[K]) => void): this;
+    once<K extends keyof ResponseEvents<T>>(type: K, listener: (event: ResponseEvents<T>[K]) => void): this;
+    bindSignal(signal: AbortSignal): this;
+    abort(reason?: string | Error): void;
+    then<TResult1 = S, TResult2 = never>(onfulfilled?: ((value: S) => TResult1 | PromiseLike<TResult1>) | null, onrejected?: ((reason: any) => TResult2 | PromiseLike<TResult2>) | null): PromiseLike<TResult1 | TResult2>;
+    catch<TResult = never>(onrejected?: ((reason: any) => TResult | PromiseLike<TResult>) | null): PromiseLike<S | TResult>;
+    result(): Promise<{
+        output: T;
+        usage: Usage;
+        elapsed: number;
+    }>;
+}
+
 type Options$6 = {
     /** The maximum number of tokens to generate */
     length?: number;
@@ -81,7 +174,7 @@ type Options$6 = {
 declare module '@botpress/zai' {
     interface Zai {
         /** Generates a text of the desired length according to the prompt */
-        text(prompt: string, options?: Options$6): Promise<string>;
+        text(prompt: string, options?: Options$6): Response<string>;
     }
 }
 
@@ -99,7 +192,7 @@ type Options$5 = {
 declare module '@botpress/zai' {
     interface Zai {
         /** Rewrites a string according to match the prompt */
-        rewrite(original: string, prompt: string, options?: Options$5): Promise<string>;
+        rewrite(original: string, prompt: string, options?: Options$5): Response<string>;
     }
 }
 
@@ -123,7 +216,7 @@ type Options$4 = {
 declare module '@botpress/zai' {
     interface Zai {
         /** Summarizes a text of any length to a summary of the desired length */
-        summarize(original: string, options?: Options$4): Promise<string>;
+        summarize(original: string, options?: Options$4): Response<string>;
     }
 }
 
@@ -140,12 +233,12 @@ type Options$3 = {
 declare module '@botpress/zai' {
     interface Zai {
         /** Checks wether a condition is true or not */
-        check(input: unknown, condition: string, options?: Options$3): Promise<{
+        check(input: unknown, condition: string, options?: Options$3): Response<{
             /** Whether the condition is true or not */
             value: boolean;
             /** The explanation of the decision */
             explanation: string;
-        }>;
+        }, boolean>;
     }
 }
 
@@ -163,7 +256,7 @@ type Options$2 = {
 declare module '@botpress/zai' {
     interface Zai {
         /** Filters elements of an array against a condition */
-        filter<T>(input: Array<T>, condition: string, options?: Options$2): Promise<Array<T>>;
+        filter<T>(input: Array<T>, condition: string, options?: Options$2): Response<Array<T>>;
     }
 }
 
@@ -182,7 +275,7 @@ type OfType<O, T extends __Z = __Z<O>> = T extends __Z<O> ? T : never;
 declare module '@botpress/zai' {
     interface Zai {
         /** Extracts one or many elements from an arbitrary input */
-        extract<S extends OfType<any>>(input: unknown, schema: S, options?: Options$1): Promise<S['_output']>;
+        extract<S extends OfType<any>>(input: unknown, schema: S, options?: Options$1): Response<S['_output']>;
     }
 }
 
@@ -213,12 +306,14 @@ type Labels<T extends string> = Record<T, string>;
 declare module '@botpress/zai' {
     interface Zai {
         /** Tags the provided input with a list of predefined labels */
-        label<T extends string>(input: unknown, labels: Labels<T>, options?: Options<T>): Promise<{
+        label<T extends string>(input: unknown, labels: Labels<T>, options?: Options<T>): Response<{
             [K in T]: {
                 explanation: string;
                 value: boolean;
                 confidence: number;
             };
+        }, {
+            [K in T]: boolean;
         }>;
     }
 }

package/dist/operations/check.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";
@@ -14,12 +17,12 @@ const _Options = z.object({
 const TRUE = "\u25A0TRUE\u25A0";
 const FALSE = "\u25A0FALSE\u25A0";
 const END = "\u25A0END\u25A0";
-Zai.prototype.check = async function(input, condition, _options) {
-  const options = _Options.parse(_options ?? {});
-  const tokenizer = await this.getTokenizer();
-  await this.fetchModelDetails();
-  const PROMPT_COMPONENT = Math.max(this.ModelDetails.input.maxTokens - PROMPT_INPUT_BUFFER, 100);
-  const taskId = this.taskId;
+const check = async (input, condition, options, ctx) => {
+  ctx.controller.signal.throwIfAborted();
+  const tokenizer = await getTokenizer();
+  const model = await ctx.getModel();
+  const PROMPT_COMPONENT = Math.max(model.input.maxTokens - PROMPT_INPUT_BUFFER, 100);
+  const taskId = ctx.taskId;
   const taskType = "zai.check";
   const PROMPT_TOKENS = {
     INPUT: Math.floor(0.5 * PROMPT_COMPONENT),
@@ -36,7 +39,7 @@ Zai.prototype.check = async function(input, condition, _options) {
       condition
     })
   );
-  const examples = taskId ? await this.adapter.getExamples({
+  const examples = taskId && ctx.adapter ? await ctx.adapter.getExamples({
     input: inputAsString,
     taskType,
     taskId
@@ -73,10 +76,10 @@ ${input2.trim()}
 <|end_input|>
 `.trim();
   };
-  const formatOutput = (answer2, justification) => {
+  const formatOutput = (answer, justification) => {
     return `
 Analysis: ${justification}
-Final Answer: ${answer2 ? TRUE : FALSE}
+Final Answer: ${answer ? TRUE : FALSE}
 ${END}
 `.trim();
   };
@@ -103,7 +106,10 @@ ${END}
 - When in doubt, ground your decision on the examples provided by the experts instead of your own intuition.
 - When no example is similar to the input, make sure to provide a clear justification for your decision while inferring the decision-making process from the examples provided by the experts.
 `.trim() : "";
-  const { output, meta } = await this.callModel({
+  const {
+    extracted: { finalAnswer, explanation },
+    meta
+  } = await ctx.generateContent({
     systemPrompt: `
 Check if the following condition is true or false for the given input. Before answering, make sure to read the input and the condition carefully.
 Justify your answer, then answer with either ${TRUE} or ${FALSE} at the very end, then add ${END} to finish the response.
@@ -123,23 +129,25 @@ ${formatInput(inputAsString, condition)}
 In your "Analysis", please refer to the Expert Examples # to justify your decision.`.trim(),
         role: "user"
       }
-    ]
+    ],
+    transform: (text) => {
+      const hasTrue = text.includes(TRUE);
+      const hasFalse = text.includes(FALSE);
+      if (!hasTrue && !hasFalse) {
+        throw new Error(`The model did not return a valid answer. The response was: ${text}`);
+      }
+      let finalAnswer2;
+      const explanation2 = text.replace(TRUE, "").replace(FALSE, "").replace(END, "").replace("Final Answer:", "").replace("Analysis:", "").trim();
+      if (hasTrue && hasFalse) {
+        finalAnswer2 = text.lastIndexOf(TRUE) > text.lastIndexOf(FALSE);
+      } else {
+        finalAnswer2 = hasTrue;
+      }
+      return { finalAnswer: finalAnswer2, explanation: explanation2.trim() };
+    }
   });
-  const answer = output.choices[0]?.content;
-  const hasTrue = answer.includes(TRUE);
-  const hasFalse = answer.includes(FALSE);
-  if (!hasTrue && !hasFalse) {
-    throw new Error(`The model did not return a valid answer. The response was: ${answer}`);
-  }
-  let finalAnswer;
-  const explanation = answer.replace(TRUE, "").replace(FALSE, "").replace(END, "").replace("Final Answer:", "").replace("Analysis:", "").trim();
-  if (hasTrue && hasFalse) {
-    finalAnswer = answer.lastIndexOf(TRUE) > answer.lastIndexOf(FALSE);
-  } else {
-    finalAnswer = hasTrue;
-  }
-  if (taskId) {
-    await this.adapter.saveExample({
+  if (taskId && ctx.adapter && !ctx.controller.signal.aborted) {
+    await ctx.adapter.saveExample({
       key: Key,
       taskType,
       taskId,
@@ -151,7 +159,7 @@ In your "Analysis", please refer to the Expert Examples # to justify your decisi
           output: meta.cost.output
         },
         latency: meta.latency,
-        model: this.Model,
+        model: ctx.modelId,
         tokens: {
           input: meta.tokens.input,
           output: meta.tokens.output
@@ -166,3 +174,14 @@ In your "Analysis", please refer to the Expert Examples # to justify your decisi
     explanation: explanation.trim()
   };
 };
+Zai.prototype.check = function(input, condition, _options) {
+  const options = _Options.parse(_options ?? {});
+  const context = new ZaiContext({
+    client: this.client,
+    modelId: this.Model,
+    taskId: this.taskId,
+    taskType: "zai.check",
+    adapter: this.adapter
+  });
+  return new Response(context, check(input, condition, options, context), (result) => result.value);
+};

package/dist/operations/extract.js

@@ -2,6 +2,9 @@ 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";
@@ -14,14 +17,15 @@ const Options = z.object({
 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;
@@ -54,10 +58,7 @@ Zai.prototype.extract = async function(input, _schema, _options) {
   }
   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) {
@@ -65,19 +66,25 @@ Zai.prototype.extract = async function(input, _schema, _options) {
     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
-        })
+        (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 this.extract(
+    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.
@@ -89,7 +96,8 @@ ${rows}
 
 Merge it back into a final result.`.trim(),
       originalSchema,
-      options
+      options,
+      ctx
     );
   }
   const instructions = [];
@@ -123,7 +131,7 @@ Merge it back into a final result.`.trim(),
       instructions: options.instructions
     })
   );
-  const examples = taskId ? await this.adapter.getExamples({
+  const examples = taskId && ctx.adapter ? await ctx.adapter.getExamples({
     input: inputAsString,
     taskType,
     taskId
@@ -182,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)}
@@ -208,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}
@@ -224,33 +232,32 @@ ${instructions.map((x) => `\u2022 ${x}`).join("\n")}
         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 && 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);
+    ],
+    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"));
       }
-      return parsedJson;
-    } catch (error) {
-      throw new JsonParsingError(x, error instanceof Error ? error : new Error("Unknown error"));
-    }
-  }).filter((x) => x !== null);
+    }).filter((x) => x !== null)
+  });
   let final;
   if (isArrayOfObjects) {
-    final = elements;
-  } else if (elements.length === 0) {
+    final = extracted;
+  } else if (extracted.length === 0) {
     final = options.strict ? schema.parse({}) : {};
   } else {
-    final = elements[0];
+    final = extracted[0];
   }
   if (wrappedValue) {
     if (Array.isArray(final)) {
@@ -259,8 +266,8 @@ ${instructions.map((x) => `\u2022 ${x}`).join("\n")}
       final = "value" in final ? final.value : final;
     }
   }
-  if (taskId) {
-    await this.adapter.saveExample({
+  if (taskId && ctx.adapter && !ctx.controller.signal.aborted) {
+    await ctx.adapter.saveExample({
       key: Key,
       taskId: `zai/${taskId}`,
       taskType,
@@ -273,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
@@ -283,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);
+};