@botpress/zai 1.1.0 → 1.2.0

package/dist/operations/extract.js

@@ -16,16 +16,22 @@ 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();
+  await this.fetchModelDetails();
   const taskId = this.taskId;
   const taskType = "zai.extract";
-  const PROMPT_COMPONENT = Math.max(this.Model.input.maxTokens - PROMPT_INPUT_BUFFER, 100);
+  const PROMPT_COMPONENT = Math.max(this.ModelDetails.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) {
+  const baseType = (schema.naked ? schema.naked() : schema)?.constructor?.name ?? "unknown";
+  if (baseType === "ZodObject") {
+  } else if (baseType === "ZodArray") {
+    let elementType = schema.element;
+    if (elementType.naked) {
+      elementType = elementType.naked();
+    }
+    if (elementType?.constructor?.name === "ZodObject") {
       isArrayOfObjects = true;
-      schema = schema._def.type;
+      schema = elementType;
     } else {
       throw new Error("Schema must be a ZodObject or a ZodArray<ZodObject>");
     }
@@ -34,7 +40,10 @@ Zai.prototype.extract = async function(input, schema, _options) {
   }
   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);
+  options.chunkLength = Math.min(
+    options.chunkLength,
+    this.ModelDetails.input.maxTokens - PROMPT_INPUT_BUFFER - schemaLength
+  );
   const keys = Object.keys(schema.shape);
   let inputAsString = stringify(input);
   if (tokenizer.count(inputAsString) > options.chunkLength) {
@@ -160,7 +169,7 @@ ${END}`.trim()
     EXAMPLES_TOKENS,
     (el) => tokenizer.count(stringify(el.input)) + tokenizer.count(stringify(el.extracted))
   ).map(formatExample).flat();
-  const output = await this.callModel({
+  const { output, meta } = await this.callModel({
     systemPrompt: `
 Extract the following information from the input:
 ${schemaTypescript}
@@ -205,7 +214,18 @@ ${instructions.map((x) => `\u2022 ${x}`).join("\n")}
       instructions: options.instructions ?? "No specific instructions",
       input: inputAsString,
       output: final,
-      metadata: output.metadata
+      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
+        }
+      }
     });
   }
   return final;

package/dist/operations/filter.js

@@ -16,10 +16,11 @@ const END = "\u25A0END\u25A0";
 Zai.prototype.filter = async function(input, condition, _options) {
   const options = Options.parse(_options ?? {});
   const tokenizer = await this.getTokenizer();
+  await this.fetchModelDetails();
   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_TOTAL_MAX = this.ModelDetails.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;
@@ -121,7 +122,7 @@ ${examples.map((x, idx) => `\u25A0${idx}:${!!x.filter ? "true" : "false"}:${x.re
         role: "assistant"
       }
     ];
-    const output = await this.callModel({
+    const { output, meta } = 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:
@@ -169,7 +170,18 @@ The condition is: "${condition}"
         input: JSON.stringify(chunk),
         output: partial,
         instructions: condition,
-        metadata: output.metadata
+        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
+          }
+        }
       });
     }
     return partial;

package/dist/operations/label.js

@@ -59,9 +59,10 @@ 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.Model.input.maxTokens - PROMPT_INPUT_BUFFER);
+  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,
@@ -152,7 +153,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 = await this.callModel({
+  const { output, meta } = await this.callModel({
     stopSequences: [END],
     systemPrompt: `
 You need to tag the input with the following labels based on the question asked:
@@ -228,7 +229,18 @@ For example, you can say: "According to Expert Example #1, ..."`.trim()
       taskType,
       taskId,
       instructions: options.instructions ?? "",
-      metadata: output.metadata,
+      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
     });

package/dist/operations/rewrite.js

@@ -15,15 +15,16 @@ 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.Model.input.maxTokens - PROMPT_INPUT_BUFFER) / 2);
+  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.Model.input.maxTokens - tokenizer.count(prompt) - PROMPT_INPUT_BUFFER;
+  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.Model.name} = ${this.Model.input.maxTokens} tokens)`
+      `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 = [];
@@ -74,13 +75,13 @@ ${END}
     ...tableExamples.map((x) => ({ input: x.input, output: x.output })),
     ...options.examples
   ];
-  const REMAINING_TOKENS = this.Model.input.maxTokens - tokenizer.count(prompt) - PROMPT_INPUT_BUFFER;
+  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 = await this.callModel({
+  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")}
@@ -99,7 +100,18 @@ ${instructions.map((x) => `\u2022 ${x}`).join("\n")}
   if (taskId) {
     await this.adapter.saveExample({
       key: Key,
-      metadata: output.metadata,
+      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,

package/dist/operations/summarize.js

@@ -20,16 +20,17 @@ const END = "\u25A0END\u25A0";
 Zai.prototype.summarize = async function(original, _options) {
   const options = Options.parse(_options ?? {});
   const tokenizer = await this.getTokenizer();
-  const INPUT_COMPONENT_SIZE = Math.max(100, (this.Model.input.maxTokens - PROMPT_INPUT_BUFFER) / 4);
+  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.Model.output.maxTokens - PROMPT_OUTPUT_BUFFER;
+  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.Model.output.maxTokens} tokens for this model (${this.Model.name})`
+      `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.Model.input.maxTokens - PROMPT_INPUT_BUFFER);
+  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 `
@@ -102,7 +103,7 @@ ${newText}
         );
       }
     }
-    const output = await this.callModel({
+    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).

package/dist/operations/text.js

@@ -8,9 +8,10 @@ const Options = z.object({
 Zai.prototype.text = async function(prompt, _options) {
   const options = Options.parse(_options ?? {});
   const tokenizer = await this.getTokenizer();
-  prompt = tokenizer.truncate(prompt, Math.max(this.Model.input.maxTokens - PROMPT_INPUT_BUFFER, 100));
+  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.Model.output.maxTokens - PROMPT_OUTPUT_BUFFER, options.length);
+    options.length = Math.min(this.ModelDetails.output.maxTokens - PROMPT_OUTPUT_BUFFER, options.length);
   }
   const instructions = [];
   let chart = "";
@@ -32,7 +33,7 @@ Zai.prototype.text = async function(prompt, _options) {
 | 200-300 tokens| A medium paragraph (150-200 words) |
 | 300-500 tokens| A long paragraph (200-300 words)   |`.trim();
   }
-  const output = await this.callModel({
+  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")}

package/dist/utils.js

@@ -2,12 +2,6 @@ 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 const BotpressClient = z.custom(
-  (value) => typeof value === "object" && value !== null && "callAction" in value && typeof value.callAction === "function",
-  {
-    message: "Invalid Botpress Client. Make sure to pass an instance of @botpress/client"
-  }
-);
 export function fastHash(str) {
   let hash = 0;
   for (let i = 0; i < str.length; i++) {

package/dist/zai.js

@@ -1,9 +1,8 @@
+import { Cognitive } from "@botpress/cognitive";
+import { getWasmTokenizer } from "@bpinternal/thicktoken";
 import { z } from "@bpinternal/zui";
-import { getWasmTokenizer } from "@botpress/wasm";
 import { TableAdapter } from "./adapters/botpress-table";
 import { MemoryAdapter } from "./adapters/memory";
-import { Models } from "./models";
-import { BotpressClient } from "./utils";
 const ActiveLearning = z.object({
   enable: z.boolean().describe("Whether to enable active learning").default(false),
   tableName: z.string().regex(
@@ -16,12 +15,14 @@ const ActiveLearning = z.object({
   ).describe("The ID of the task").default("default")
 });
 const ZaiConfig = z.object({
-  client: BotpressClient,
+  client: z.custom(),
   userId: z.string().describe("The ID of the user consuming the API").optional(),
-  retry: z.object({ maxRetries: z.number().min(0).max(100) }).default({ maxRetries: 3 }),
   modelId: z.custom(
     (value) => {
-      if (typeof value !== "string" || !value.includes("__")) {
+      if (typeof value !== "string") {
+        return false;
+      }
+      if (value !== "best" && value !== "fast" && !value.includes(":")) {
         return false;
       }
       return true;
@@ -29,7 +30,7 @@ const ZaiConfig = z.object({
     {
       message: "Invalid model ID"
     }
-  ).describe("The ID of the model you want to use").default("openai__gpt-4o-mini-2024-07-18"),
+  ).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}$/,
@@ -39,76 +40,30 @@ const ZaiConfig = z.object({
 export class Zai {
   static tokenizer = null;
   client;
-  originalConfig;
-  userId;
-  integration;
-  model;
-  retry;
+  _originalConfig;
+  _userId;
   Model;
+  ModelDetails;
   namespace;
   adapter;
   activeLearning;
   constructor(config) {
-    this.originalConfig = config;
+    this._originalConfig = config;
     const parsed = ZaiConfig.parse(config);
-    this.client = parsed.client;
-    const [integration, modelId] = parsed.modelId.split("__");
-    if (!integration?.length || !modelId?.length) {
-      throw new Error(`Invalid model ID: ${parsed.modelId}. Expected format: <integration>__<modelId>`);
-    }
-    this.integration = integration;
-    this.model = modelId;
+    this.client = Cognitive.isCognitiveClient(parsed.client) ? parsed.client : new Cognitive({ client: parsed.client });
     this.namespace = parsed.namespace;
-    this.userId = parsed.userId;
-    this.retry = parsed.retry;
-    this.Model = Models.find((m) => m.id === parsed.modelId);
+    this._userId = parsed.userId;
+    this.Model = parsed.modelId;
     this.activeLearning = parsed.activeLearning;
-    this.adapter = parsed.activeLearning?.enable ? new TableAdapter({ client: this.client, tableName: parsed.activeLearning.tableName }) : new MemoryAdapter([]);
+    this.adapter = parsed.activeLearning?.enable ? new TableAdapter({ client: this.client.client, tableName: parsed.activeLearning.tableName }) : new MemoryAdapter([]);
   }
   /** @internal */
   async callModel(props) {
-    let retries = this.retry.maxRetries;
-    while (retries-- >= 0) {
-      try {
-        return await this._callModel(props);
-      } catch (e) {
-        if (retries >= 0) {
-          await new Promise((resolve) => setTimeout(resolve, 1e3));
-        } else {
-          throw new Error("Failed to call model after multiple retries");
-        }
-      }
-    }
-    throw new Error("Failed to call model after multiple retries");
-  }
-  /** @internal */
-  async _callModel(props) {
-    let retries = this.retry.maxRetries;
-    do {
-      const start = Date.now();
-      const input = {
-        messages: [],
-        temperature: 0,
-        topP: 1,
-        model: { id: this.model },
-        userId: this.userId,
-        ...props
-      };
-      const { output } = await this.client.callAction({
-        type: `${this.integration}:generateContent`,
-        input
-      });
-      const latency = Date.now() - start;
-      return {
-        ...output,
-        metadata: {
-          model: this.model,
-          latency,
-          cost: { input: output.usage.inputCost, output: output.usage.outputCost },
-          tokens: { input: output.usage.inputTokens, output: output.usage.outputTokens }
-        }
-      };
-    } while (--retries > 0);
+    return this.client.generateContent({
+      ...props,
+      model: this.Model,
+      userId: this._userId
+    });
   }
   async getTokenizer() {
     Zai.tokenizer ??= await (async () => {
@@ -119,6 +74,11 @@ export class Zai {
     })();
     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;
@@ -127,13 +87,13 @@ export class Zai {
   }
   with(options) {
     return new Zai({
-      ...this.originalConfig,
+      ...this._originalConfig,
       ...options
     });
   }
   learn(taskId) {
     return new Zai({
-      ...this.originalConfig,
+      ...this._originalConfig,
       activeLearning: { ...this.activeLearning, taskId, enable: true }
     });
   }