@botpress/zai 2.0.6 → 2.0.8

package/dist/adapters/botpress-table.js

@@ -1,5 +1,4 @@
 import { z } from "@bpinternal/zui";
-import { GenerationMetadata } from "../utils";
 import { Adapter } from "./adapter";
 const CRITICAL_TAGS = {
   system: "true",
@@ -29,7 +28,7 @@ const TableSchema = z.object({
   input: z.object({}).passthrough().describe("The input to the task"),
   output: z.object({}).passthrough().describe("The expected output"),
   explanation: z.string().nullable(),
-  metadata: GenerationMetadata,
+  metadata: z.object({}).passthrough(),
   status: z.enum(["pending", "rejected", "approved"]),
   feedback: z.object({
     rating: z.enum(["very-bad", "bad", "good", "very-good"]),
@@ -105,7 +104,9 @@ export class TableAdapter extends Adapter {
           output: { value: output },
           explanation: explanation ?? null,
           status,
-          metadata
+          metadata,
+          feedback: null
+          // Feedback is not provided at this point
         }
       ]
     }).catch(() => {

package/dist/index.d.ts

@@ -1,55 +1,18 @@
-import * as _bpinternal_zui from '@bpinternal/zui';
-import { z, ZodObject } from '@bpinternal/zui';
 import { Cognitive, Model, BotpressClientLike } from '@botpress/cognitive';
 import { TextTokenizer } from '@bpinternal/thicktoken';
 
-type GenerationMetadata = (typeof GenerationMetadata)['_input'];
-declare const GenerationMetadata: _bpinternal_zui.ZodObject<{
-    model: _bpinternal_zui.ZodString;
-    cost: _bpinternal_zui.ZodObject<{
-        input: _bpinternal_zui.ZodNumber;
-        output: _bpinternal_zui.ZodNumber;
-    }, "strip", _bpinternal_zui.ZodTypeAny, {
-        input?: number;
-        output?: number;
-    }, {
-        input?: number;
-        output?: number;
-    }>;
-    latency: _bpinternal_zui.ZodNumber;
-    tokens: _bpinternal_zui.ZodObject<{
-        input: _bpinternal_zui.ZodNumber;
-        output: _bpinternal_zui.ZodNumber;
-    }, "strip", _bpinternal_zui.ZodTypeAny, {
-        input?: number;
-        output?: number;
-    }, {
-        input?: number;
-        output?: number;
-    }>;
-}, "strip", _bpinternal_zui.ZodTypeAny, {
-    model?: string;
-    cost?: {
-        input?: number;
-        output?: number;
-    };
-    latency?: number;
-    tokens?: {
-        input?: number;
-        output?: number;
-    };
-}, {
-    model?: string;
-    cost?: {
-        input?: number;
-        output?: number;
+type GenerationMetadata = {
+    model: string;
+    cost: {
+        input: number;
+        output: number;
     };
-    latency?: number;
-    tokens?: {
-        input?: number;
-        output?: number;
+    latency: number;
+    tokens: {
+        input: number;
+        output: number;
     };
-}>;
+};
 
 type SaveExampleProps<TInput, TOutput> = {
     key: string;
@@ -79,60 +42,18 @@ declare abstract class Adapter {
 }
 
 type ModelId = Required<Parameters<Cognitive['generateContent']>[0]['model']>;
-type ActiveLearning = (typeof ActiveLearning)['_input'];
-declare const ActiveLearning: _bpinternal_zui.ZodObject<{
-    enable: _bpinternal_zui.ZodDefault<_bpinternal_zui.ZodBoolean>;
-    tableName: _bpinternal_zui.ZodDefault<_bpinternal_zui.ZodString>;
-    taskId: _bpinternal_zui.ZodDefault<_bpinternal_zui.ZodString>;
-}, "strip", _bpinternal_zui.ZodTypeAny, {
-    tableName?: string;
-    taskId?: string;
-    enable?: boolean;
-}, {
-    tableName?: string;
-    taskId?: string;
-    enable?: boolean;
-}>;
-type ZaiConfig = (typeof ZaiConfig)['_input'];
-declare const ZaiConfig: _bpinternal_zui.ZodObject<{
-    client: z.Schema<Cognitive | BotpressClientLike, _bpinternal_zui.ZodTypeDef, Cognitive | BotpressClientLike>;
-    userId: _bpinternal_zui.ZodOptional<_bpinternal_zui.ZodString>;
-    modelId: _bpinternal_zui.ZodDefault<z.Schema<string, _bpinternal_zui.ZodTypeDef, string>>;
-    activeLearning: _bpinternal_zui.ZodDefault<_bpinternal_zui.ZodObject<{
-        enable: _bpinternal_zui.ZodDefault<_bpinternal_zui.ZodBoolean>;
-        tableName: _bpinternal_zui.ZodDefault<_bpinternal_zui.ZodString>;
-        taskId: _bpinternal_zui.ZodDefault<_bpinternal_zui.ZodString>;
-    }, "strip", _bpinternal_zui.ZodTypeAny, {
-        tableName?: string;
-        taskId?: string;
-        enable?: boolean;
-    }, {
-        tableName?: string;
-        taskId?: string;
-        enable?: boolean;
-    }>>;
-    namespace: _bpinternal_zui.ZodDefault<_bpinternal_zui.ZodString>;
-}, "strip", _bpinternal_zui.ZodTypeAny, {
-    client?: Cognitive | BotpressClientLike;
-    userId?: string;
-    modelId?: string;
-    activeLearning?: {
-        tableName?: string;
-        taskId?: string;
-        enable?: boolean;
-    };
-    namespace?: string;
-}, {
-    client?: Cognitive | BotpressClientLike;
+type ActiveLearning = {
+    enable: boolean;
+    tableName: string;
+    taskId: string;
+};
+type ZaiConfig = {
+    client: BotpressClientLike | Cognitive;
     userId?: string;
-    modelId?: string;
-    activeLearning?: {
-        tableName?: string;
-        taskId?: string;
-        enable?: boolean;
-    };
+    modelId?: ModelId | string;
+    activeLearning?: ActiveLearning;
     namespace?: string;
-}>;
+};
 declare class Zai {
     protected static tokenizer: TextTokenizer;
     protected client: Cognitive;
@@ -153,14 +74,10 @@ declare class Zai {
     learn(taskId: string): Zai;
 }
 
-type Options$6 = (typeof Options$6)['_input'];
-declare const Options$6: _bpinternal_zui.ZodObject<{
-    length: _bpinternal_zui.ZodOptional<_bpinternal_zui.ZodNumber>;
-}, "strip", _bpinternal_zui.ZodTypeAny, {
+type Options$6 = {
+    /** The maximum number of tokens to generate */
     length?: number;
-}, {
-    length?: number;
-}>;
+};
 declare module '@botpress/zai' {
     interface Zai {
         /** Generates a text of the desired length according to the prompt */
@@ -168,32 +85,17 @@ declare module '@botpress/zai' {
     }
 }
 
-type Options$5 = (typeof Options$5)['_input'];
-declare const Options$5: _bpinternal_zui.ZodObject<{
-    examples: _bpinternal_zui.ZodDefault<_bpinternal_zui.ZodArray<_bpinternal_zui.ZodObject<{
-        input: _bpinternal_zui.ZodString;
-        output: _bpinternal_zui.ZodString;
-    }, "strip", _bpinternal_zui.ZodTypeAny, {
-        input?: string;
-        output?: string;
-    }, {
-        input?: string;
-        output?: string;
-    }>, "many">>;
-    length: _bpinternal_zui.ZodOptional<_bpinternal_zui.ZodNumber>;
-}, "strip", _bpinternal_zui.ZodTypeAny, {
-    length?: number;
-    examples?: {
-        input?: string;
-        output?: string;
-    }[];
-}, {
+type Example$2 = {
+    input: string;
+    output: string;
+    instructions?: string;
+};
+type Options$5 = {
+    /** Examples to guide the rewriting */
+    examples?: Array<Example$2>;
+    /** The maximum number of tokens to generate */
     length?: number;
-    examples?: {
-        input?: string;
-        output?: string;
-    }[];
-}>;
+};
 declare module '@botpress/zai' {
     interface Zai {
         /** Rewrites a string according to match the prompt */
@@ -201,44 +103,23 @@ declare module '@botpress/zai' {
     }
 }
 
-type Options$4 = (typeof Options$4)['_input'];
-declare const Options$4: _bpinternal_zui.ZodObject<{
-    prompt: _bpinternal_zui.ZodDefault<_bpinternal_zui.ZodString>;
-    format: _bpinternal_zui.ZodDefault<_bpinternal_zui.ZodString>;
-    length: _bpinternal_zui.ZodDefault<_bpinternal_zui.ZodNumber>;
-    intermediateFactor: _bpinternal_zui.ZodDefault<_bpinternal_zui.ZodNumber>;
-    maxIterations: _bpinternal_zui.ZodDefault<_bpinternal_zui.ZodNumber>;
-    sliding: _bpinternal_zui.ZodDefault<_bpinternal_zui.ZodObject<{
-        window: _bpinternal_zui.ZodNumber;
-        overlap: _bpinternal_zui.ZodNumber;
-    }, "strip", _bpinternal_zui.ZodTypeAny, {
-        window?: number;
-        overlap?: number;
-    }, {
-        window?: number;
-        overlap?: number;
-    }>>;
-}, "strip", _bpinternal_zui.ZodTypeAny, {
-    length?: number;
+type Options$4 = {
+    /** What should the text be summarized to? */
     prompt?: string;
+    /** How to format the example text */
     format?: string;
-    intermediateFactor?: number;
-    maxIterations?: number;
-    sliding?: {
-        window?: number;
-        overlap?: number;
-    };
-}, {
+    /** The length of the summary in tokens */
     length?: number;
-    prompt?: string;
-    format?: string;
+    /** How many times longer (than final length) are the intermediate summaries generated */
     intermediateFactor?: number;
+    /** The maximum number of iterations to perform */
     maxIterations?: number;
+    /** Sliding window options */
     sliding?: {
-        window?: number;
-        overlap?: number;
+        window: number;
+        overlap: number;
     };
-}>;
+};
 declare module '@botpress/zai' {
     interface Zai {
         /** Summarizes a text of any length to a summary of the desired length */
@@ -246,34 +127,14 @@ declare module '@botpress/zai' {
     }
 }
 
-type Options$3 = (typeof Options$3)['_input'];
-declare const Options$3: _bpinternal_zui.ZodObject<{
-    examples: _bpinternal_zui.ZodDefault<_bpinternal_zui.ZodArray<_bpinternal_zui.ZodObject<{
-        input: _bpinternal_zui.ZodAny;
-        check: _bpinternal_zui.ZodBoolean;
-        reason: _bpinternal_zui.ZodOptional<_bpinternal_zui.ZodString>;
-    }, "strip", _bpinternal_zui.ZodTypeAny, {
-        input?: any;
-        check?: boolean;
-        reason?: string;
-    }, {
-        input?: any;
-        check?: boolean;
-        reason?: string;
-    }>, "many">>;
-}, "strip", _bpinternal_zui.ZodTypeAny, {
-    examples?: {
-        input?: any;
-        check?: boolean;
+type Options$3 = {
+    /** Examples to check the condition against */
+    examples?: Array<{
+        input: unknown;
+        check: boolean;
         reason?: string;
-    }[];
-}, {
-    examples?: {
-        input?: any;
-        check?: boolean;
-        reason?: string;
-    }[];
-}>;
+    }>;
+};
 declare module '@botpress/zai' {
     interface Zai {
         /** Checks wether a condition is true or not */
@@ -286,37 +147,17 @@ declare module '@botpress/zai' {
     }
 }
 
-type Options$2 = (typeof Options$2)['_input'];
-declare const Options$2: _bpinternal_zui.ZodObject<{
-    tokensPerItem: _bpinternal_zui.ZodDefault<_bpinternal_zui.ZodOptional<_bpinternal_zui.ZodNumber>>;
-    examples: _bpinternal_zui.ZodDefault<_bpinternal_zui.ZodArray<_bpinternal_zui.ZodObject<{
-        input: _bpinternal_zui.ZodAny;
-        filter: _bpinternal_zui.ZodBoolean;
-        reason: _bpinternal_zui.ZodOptional<_bpinternal_zui.ZodString>;
-    }, "strip", _bpinternal_zui.ZodTypeAny, {
-        input?: any;
-        filter?: boolean;
-        reason?: string;
-    }, {
-        input?: any;
-        filter?: boolean;
-        reason?: string;
-    }>, "many">>;
-}, "strip", _bpinternal_zui.ZodTypeAny, {
-    examples?: {
-        input?: any;
-        filter?: boolean;
-        reason?: string;
-    }[];
-    tokensPerItem?: number;
-}, {
-    examples?: {
-        input?: any;
-        filter?: boolean;
-        reason?: string;
-    }[];
+type Example$1 = {
+    input: unknown;
+    filter: boolean;
+    reason?: string;
+};
+type Options$2 = {
+    /** The maximum number of tokens per item */
     tokensPerItem?: number;
-}>;
+    /** Examples to filter the condition against */
+    examples?: Array<Example$1>;
+};
 declare module '@botpress/zai' {
     interface Zai {
         /** Filters elements of an array against a condition */
@@ -324,17 +165,12 @@ declare module '@botpress/zai' {
     }
 }
 
-type Options$1 = (typeof Options$1)['_input'];
-declare const Options$1: ZodObject<{
-    instructions: _bpinternal_zui.ZodOptional<_bpinternal_zui.ZodString>;
-    chunkLength: _bpinternal_zui.ZodDefault<_bpinternal_zui.ZodOptional<_bpinternal_zui.ZodNumber>>;
-}, "strip", _bpinternal_zui.ZodTypeAny, {
-    instructions?: string;
-    chunkLength?: number;
-}, {
+type Options$1 = {
+    /** Instructions to guide the user on how to extract the data */
     instructions?: string;
+    /** The maximum number of tokens per chunk */
     chunkLength?: number;
-}>;
+};
 type __Z<T extends any = any> = {
     _output: T;
 };
@@ -362,60 +198,15 @@ type Example<T extends string> = {
         explanation?: string;
     }>>;
 };
-type Options<T extends string> = Omit<(typeof Options)['_input'], 'examples'> & {
-    examples?: Array<Partial<Example<T>>>;
-};
-declare const Options: _bpinternal_zui.ZodObject<{
-    examples: _bpinternal_zui.ZodDefault<_bpinternal_zui.ZodArray<_bpinternal_zui.ZodObject<{
-        input: _bpinternal_zui.ZodAny;
-        labels: _bpinternal_zui.ZodRecord<_bpinternal_zui.ZodString, _bpinternal_zui.ZodObject<{
-            label: _bpinternal_zui.ZodEnum<never>;
-            explanation: _bpinternal_zui.ZodOptional<_bpinternal_zui.ZodString>;
-        }, "strip", _bpinternal_zui.ZodTypeAny, {
-            label: never;
-            explanation?: string;
-        }, {
-            label: never;
-            explanation?: string;
-        }>>;
-    }, "strip", _bpinternal_zui.ZodTypeAny, {
-        input?: any;
-        labels?: Record<string, {
-            label: never;
-            explanation?: string;
-        }>;
-    }, {
-        input?: any;
-        labels?: Record<string, {
-            label: never;
-            explanation?: string;
-        }>;
-    }>, "many">>;
-    instructions: _bpinternal_zui.ZodOptional<_bpinternal_zui.ZodString>;
-    chunkLength: _bpinternal_zui.ZodDefault<_bpinternal_zui.ZodOptional<_bpinternal_zui.ZodNumber>>;
-}, "strip", _bpinternal_zui.ZodTypeAny, {
-    instructions?: string;
-    examples?: {
-        input?: any;
-        labels?: Record<string, {
-            label: never;
-            explanation?: string;
-        }>;
-    }[];
-    chunkLength?: number;
-}, {
+type Options<T extends string> = {
+    /** Examples to help the user make a decision */
+    examples?: Array<Example<T>>;
+    /** Instructions to guide the user on how to extract the data */
     instructions?: string;
-    examples?: {
-        input?: any;
-        labels?: Record<string, {
-            label: never;
-            explanation?: string;
-        }>;
-    }[];
+    /** The maximum number of tokens per chunk */
     chunkLength?: number;
-}>;
+};
 type Labels<T extends string> = Record<T, string>;
-declare const Labels: z.ZodTransformer<_bpinternal_zui.ZodRecord<_bpinternal_zui.ZodString, _bpinternal_zui.ZodString>, Record<string, string>, Record<string, string>>;
 declare module '@botpress/zai' {
     interface Zai {
         /** Tags the provided input with a list of predefined labels */

package/dist/operations/check.js

@@ -2,19 +2,19 @@ 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({
+const _Example = z.object({
   input: z.any(),
   check: z.boolean(),
   reason: z.string().optional()
 });
-const Options = z.object({
-  examples: z.array(Example).describe("Examples to check the condition against").default([])
+const _Options = z.object({
+  examples: z.array(_Example).describe("Examples to check the condition against").default([])
 });
 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 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);

package/dist/operations/filter.js

@@ -3,18 +3,18 @@ import { clamp } from "lodash-es";
 import { fastHash, stringify, takeUntilTokens } from "../utils";
 import { Zai } from "../zai";
 import { PROMPT_INPUT_BUFFER, PROMPT_OUTPUT_BUFFER } from "./constants";
-const Example = z.object({
+const _Example = z.object({
   input: z.any(),
   filter: z.boolean(),
   reason: z.string().optional()
 });
-const Options = z.object({
+const _Options = z.object({
   tokensPerItem: z.number().min(1).max(1e5).optional().describe("The maximum number of tokens per item").default(250),
-  examples: z.array(Example).describe("Examples to filter the condition against").default([])
+  examples: z.array(_Example).describe("Examples to filter the condition against").default([])
 });
 const END = "\u25A0END\u25A0";
 Zai.prototype.filter = async function(input, condition, _options) {
-  const options = Options.parse(_options ?? {});
+  const options = _Options.parse(_options ?? {});
   const tokenizer = await this.getTokenizer();
   await this.fetchModelDetails();
   const taskId = this.taskId;

package/dist/operations/label.js

@@ -11,7 +11,7 @@ const LABELS = {
   ABSOLUTELY_YES: "ABSOLUTELY_YES"
 };
 const ALL_LABELS = Object.values(LABELS).join(" | ");
-const Options = z.object({
+const _Options = z.object({
   examples: z.array(
     z.object({
       input: z.any(),
@@ -21,7 +21,7 @@ const Options = z.object({
   instructions: z.string().optional().describe("Instructions to guide the user on how to extract the data"),
   chunkLength: z.number().min(100).max(1e5).optional().describe("The maximum number of tokens per chunk").default(16e3)
 });
-const Labels = z.record(z.string().min(1).max(250), z.string()).superRefine((labels, ctx) => {
+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) {
@@ -68,8 +68,8 @@ const getConfidence = (label) => {
   }
 };
 Zai.prototype.label = async function(input, _labels, _options) {
-  const options = Options.parse(_options ?? {});
-  const labels = Labels.parse(_labels);
+  const options = _Options.parse(_options ?? {});
+  const labels = _Labels.parse(_labels);
   const tokenizer = await this.getTokenizer();
   await this.fetchModelDetails();
   const taskId = this.taskId;
@@ -126,7 +126,7 @@ Zai.prototype.label = async function(input, _labels, _options) {
   options.examples.forEach((example) => {
     examples.push({
       key: fastHash(JSON.stringify(example)),
-      input: example.input,
+      input: stringify(example.input),
       similarity: 1,
       explanation: "",
       output: example.labels

package/dist/operations/rewrite.js

@@ -2,12 +2,12 @@ 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({
+const _Example = z.object({
   input: z.string(),
   output: z.string()
 });
 const Options = z.object({
-  examples: z.array(Example).default([]),
+  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";

package/dist/utils.js

@@ -1,4 +1,3 @@
-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>";
 };
@@ -23,15 +22,3 @@ export const takeUntilTokens = (arr, tokens, count) => {
   }
   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

@@ -3,7 +3,7 @@ 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({
+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$/,
@@ -14,7 +14,7 @@ const ActiveLearning = z.object({
     "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({
+const _ZaiConfig = z.object({
   client: z.custom(),
   userId: z.string().describe("The ID of the user consuming the API").optional(),
   modelId: z.custom(
@@ -31,7 +31,7 @@ const ZaiConfig = z.object({
       message: "Invalid model ID"
     }
   ).describe("The ID of the model you want to use").default("best"),
-  activeLearning: ActiveLearning.default({ enable: false }),
+  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"
@@ -49,13 +49,16 @@ export class Zai {
   activeLearning;
   constructor(config) {
     this._originalConfig = config;
-    const parsed = ZaiConfig.parse(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([]);
+    this.adapter = parsed.activeLearning?.enable ? new TableAdapter({
+      client: this.client.client,
+      tableName: parsed.activeLearning.tableName
+    }) : new MemoryAdapter([]);
   }
   /** @internal */
   async callModel(props) {

package/package.json

@@ -1,7 +1,7 @@
 {
   "name": "@botpress/zai",
   "description": "Zui AI (zai) – An LLM utility library written on top of Zui and the Botpress API",
-  "version": "2.0.6",
+  "version": "2.0.8",
   "main": "./dist/index.js",
   "types": "./dist/index.d.ts",
   "exports": {
@@ -23,7 +23,7 @@
   "author": "",
   "license": "ISC",
   "dependencies": {
-    "@botpress/cognitive": "0.1.21",
+    "@botpress/cognitive": "0.1.23",
     "json5": "^2.2.3",
     "jsonrepair": "^3.10.0",
     "lodash-es": "^4.17.21"
@@ -41,7 +41,7 @@
   },
   "peerDependencies": {
     "@bpinternal/thicktoken": "^1.0.0",
-    "@bpinternal/zui": "^0.17.1"
+    "@bpinternal/zui": "^0.22.5"
   },
   "engines": {
     "node": ">=18.0.0"

package/src/adapters/botpress-table.ts

@@ -20,6 +20,11 @@ const OPTIONAL_TAGS = {
 
 const FACTOR = 30
 
+type Props = {
+  client: Client
+  tableName: string
+}
+
 const Props = z.object({
   client: z.custom(() => true),
   tableName: z
@@ -30,7 +35,22 @@ const Props = z.object({
     ),
 })
 
-export type TableSchema = (typeof TableSchema)['_input']
+export type TableSchema = {
+  taskType: string
+  taskId: string
+  key: string
+  instructions: string
+  input: Record<string, unknown>
+  output: Record<string, unknown>
+  explanation: string | null
+  metadata: GenerationMetadata
+  status: 'pending' | 'rejected' | 'approved'
+  feedback: {
+    rating: 'very-bad' | 'bad' | 'good' | 'very-good'
+    comment: string | null
+  } | null
+}
+
 const TableSchema = z.object({
   taskType: z.string().describe('The type of the task (filter, extract, etc.)'),
   taskId: z.string(),
@@ -39,7 +59,7 @@ const TableSchema = z.object({
   input: z.object({}).passthrough().describe('The input to the task'),
   output: z.object({}).passthrough().describe('The expected output'),
   explanation: z.string().nullable(),
-  metadata: GenerationMetadata,
+  metadata: z.object({}).passthrough(),
   status: z.enum(['pending', 'rejected', 'approved']),
   feedback: z
     .object({
@@ -65,9 +85,9 @@ export class TableAdapter extends Adapter {
 
   private _status: 'initialized' | 'ready' | 'error'
 
-  public constructor(props: (typeof Props)['_input']) {
+  public constructor(props: Props) {
     super()
-    props = Props.parse(props)
+    props = Props.parse(props) as Props
     this._client = props.client
     this._tableName = props.tableName
     this._status = 'ready'
@@ -131,6 +151,7 @@ export class TableAdapter extends Adapter {
             explanation: explanation ?? null,
             status,
             metadata,
+            feedback: null, // Feedback is not provided at this point
           } satisfies TableSchema,
         ],
       })

package/src/operations/check.ts

@@ -5,15 +5,23 @@ import { fastHash, stringify, takeUntilTokens } from '../utils'
 import { Zai } from '../zai'
 import { PROMPT_INPUT_BUFFER } from './constants'
 
-const Example = z.object({
+const _Example = z.object({
   input: z.any(),
   check: z.boolean(),
   reason: z.string().optional(),
 })
 
-export type Options = (typeof Options)['_input']
-const Options = z.object({
-  examples: z.array(Example).describe('Examples to check the condition against').default([]),
+export type Options = {
+  /** Examples to check the condition against */
+  examples?: Array<{
+    input: unknown
+    check: boolean
+    reason?: string
+  }>
+}
+
+const _Options = z.object({
+  examples: z.array(_Example).describe('Examples to check the condition against').default([]),
 })
 
 declare module '@botpress/zai' {
@@ -36,8 +44,8 @@ const TRUE = '■TRUE■'
 const FALSE = '■FALSE■'
 const END = '■END■'
 
-Zai.prototype.check = async function (this: Zai, input, condition, _options) {
-  const options = Options.parse(_options ?? {})
+Zai.prototype.check = async function (this: Zai, input: unknown, condition: string, _options: Options | undefined) {
+  const options = _Options.parse(_options ?? {}) as Options
   const tokenizer = await this.getTokenizer()
   await this.fetchModelDetails()
   const PROMPT_COMPONENT = Math.max(this.ModelDetails.input.maxTokens - PROMPT_INPUT_BUFFER, 100)

package/src/operations/extract.ts

@@ -10,7 +10,13 @@ import { Zai } from '../zai'
 import { PROMPT_INPUT_BUFFER } from './constants'
 import { JsonParsingError } from './errors'
 
-export type Options = (typeof Options)['_input']
+export type Options = {
+  /** Instructions to guide the user on how to extract the data */
+  instructions?: string
+  /** The maximum number of tokens per chunk */
+  chunkLength?: number
+}
+
 const Options = z.object({
   instructions: z.string().optional().describe('Instructions to guide the user on how to extract the data'),
   chunkLength: z

package/src/operations/filter.ts

@@ -6,15 +6,26 @@ import { fastHash, stringify, takeUntilTokens } from '../utils'
 import { Zai } from '../zai'
 import { PROMPT_INPUT_BUFFER, PROMPT_OUTPUT_BUFFER } from './constants'
 
-type Example = (typeof Example)['_input']
-const Example = z.object({
+type Example = {
+  input: unknown
+  filter: boolean
+  reason?: string
+}
+
+const _Example = z.object({
   input: z.any(),
   filter: z.boolean(),
   reason: z.string().optional(),
 })
 
-export type Options = (typeof Options)['_input']
-const Options = z.object({
+export type Options = {
+  /** The maximum number of tokens per item */
+  tokensPerItem?: number
+  /** Examples to filter the condition against */
+  examples?: Array<Example>
+}
+
+const _Options = z.object({
   tokensPerItem: z
     .number()
     .min(1)
@@ -22,7 +33,7 @@ const Options = z.object({
     .optional()
     .describe('The maximum number of tokens per item')
     .default(250),
-  examples: z.array(Example).describe('Examples to filter the condition against').default([]),
+  examples: z.array(_Example).describe('Examples to filter the condition against').default([]),
 })
 
 declare module '@botpress/zai' {
@@ -35,7 +46,7 @@ declare module '@botpress/zai' {
 const END = '■END■'
 
 Zai.prototype.filter = async function (this: Zai, input, condition, _options) {
-  const options = Options.parse(_options ?? {})
+  const options = _Options.parse(_options ?? {}) as Options
   const tokenizer = await this.getTokenizer()
   await this.fetchModelDetails()
 

package/src/operations/label.ts

@@ -22,11 +22,16 @@ type Example<T extends string> = {
   labels: Partial<Record<T, { label: Label; explanation?: string }>>
 }
 
-export type Options<T extends string> = Omit<(typeof Options)['_input'], 'examples'> & {
-  examples?: Array<Partial<Example<T>>>
+export type Options<T extends string> = {
+  /** Examples to help the user make a decision */
+  examples?: Array<Example<T>>
+  /** Instructions to guide the user on how to extract the data */
+  instructions?: string
+  /** The maximum number of tokens per chunk */
+  chunkLength?: number
 }
 
-const Options = z.object({
+const _Options = z.object({
   examples: z
     .array(
       z.object({
@@ -48,7 +53,7 @@ const Options = z.object({
 
 type Labels<T extends string> = Record<T, string>
 
-const Labels = z.record(z.string().min(1).max(250), z.string()).superRefine((labels, ctx) => {
+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) {
@@ -119,9 +124,14 @@ const getConfidence = (label: Label) => {
   }
 }
 
-Zai.prototype.label = async function <T extends string>(this: Zai, input, _labels, _options) {
-  const options = Options.parse(_options ?? {})
-  const labels = Labels.parse(_labels)
+Zai.prototype.label = async function <T extends string>(
+  this: Zai,
+  input: unknown,
+  _labels: Labels<T>,
+  _options: Options<T> | undefined
+) {
+  const options = _Options.parse(_options ?? {}) as unknown as Options<T>
+  const labels = _Labels.parse(_labels) as Labels<T>
   const tokenizer = await this.getTokenizer()
   await this.fetchModelDetails()
 
@@ -211,7 +221,7 @@ Zai.prototype.label = async function <T extends string>(this: Zai, input, _label
   options.examples.forEach((example) => {
     examples.push({
       key: fastHash(JSON.stringify(example)),
-      input: example.input,
+      input: stringify(example.input),
       similarity: 1,
       explanation: '',
       output: example.labels as unknown as {

package/src/operations/rewrite.ts

@@ -5,15 +5,26 @@ import { fastHash, stringify, takeUntilTokens } from '../utils'
 import { Zai } from '../zai'
 import { PROMPT_INPUT_BUFFER } from './constants'
 
-type Example = (typeof Example)['_input'] & { instructions?: string }
-const Example = z.object({
+type Example = {
+  input: string
+  output: string
+  instructions?: string
+}
+
+const _Example = z.object({
   input: z.string(),
   output: z.string(),
 })
 
-export type Options = (typeof Options)['_input']
+export type Options = {
+  /** Examples to guide the rewriting */
+  examples?: Array<Example>
+  /** The maximum number of tokens to generate */
+  length?: number
+}
+
 const Options = z.object({
-  examples: z.array(Example).default([]),
+  examples: z.array(_Example).default([]),
   length: z.number().min(10).max(16_000).optional().describe('The maximum number of tokens to generate'),
 })
 
@@ -28,7 +39,7 @@ const START = '■START■'
 const END = '■END■'
 
 Zai.prototype.rewrite = async function (this: Zai, original, prompt, _options) {
-  const options = Options.parse(_options ?? {})
+  const options = Options.parse(_options ?? {}) as Options
   const tokenizer = await this.getTokenizer()
   await this.fetchModelDetails()
 
@@ -101,7 +112,7 @@ ${END}
   }
 
   const savedExamples: Example[] = [
-    ...tableExamples.map((x) => ({ input: x.input as string, output: x.output as string })),
+    ...tableExamples.map((x) => ({ input: x.input as string, output: x.output as string }) satisfies Example),
     ...options.examples,
   ]
 

package/src/operations/summarize.ts

@@ -5,7 +5,24 @@ import { chunk } from 'lodash-es'
 import { Zai } from '../zai'
 import { PROMPT_INPUT_BUFFER, PROMPT_OUTPUT_BUFFER } from './constants'
 
-export type Options = (typeof Options)['_input']
+export type Options = {
+  /** What should the text be summarized to? */
+  prompt?: string
+  /** How to format the example text */
+  format?: string
+  /** The length of the summary in tokens */
+  length?: number
+  /** How many times longer (than final length) are the intermediate summaries generated */
+  intermediateFactor?: number
+  /** The maximum number of iterations to perform */
+  maxIterations?: number
+  /** Sliding window options */
+  sliding?: {
+    window: number
+    overlap: number
+  }
+}
+
 const Options = z.object({
   prompt: z
     .string()
@@ -45,7 +62,7 @@ const START = '■START■'
 const END = '■END■'
 
 Zai.prototype.summarize = async function (this: Zai, original, _options) {
-  const options = Options.parse(_options ?? {})
+  const options = Options.parse(_options ?? {}) as Options
   const tokenizer = await this.getTokenizer()
   await this.fetchModelDetails()
 

package/src/operations/text.ts

@@ -5,7 +5,11 @@ import { clamp } from 'lodash-es'
 import { Zai } from '../zai'
 import { PROMPT_INPUT_BUFFER, PROMPT_OUTPUT_BUFFER } from './constants'
 
-export type Options = (typeof Options)['_input']
+export type Options = {
+  /** The maximum number of tokens to generate */
+  length?: number
+}
+
 const Options = z.object({
   length: z.number().min(1).max(100_000).optional().describe('The maximum number of tokens to generate'),
 })

package/src/utils.ts

@@ -1,5 +1,3 @@
-import { z } from '@bpinternal/zui'
-
 export const stringify = (input: unknown, beautify = true) => {
   return typeof input === 'string' && !!input.length
     ? input
@@ -33,20 +31,15 @@ export const takeUntilTokens = <T>(arr: T[], tokens: number, count: (el: T) => n
   return result
 }
 
-export type GenerationMetadata = (typeof GenerationMetadata)['_input']
-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'),
-})
+export type GenerationMetadata = {
+  model: string
+  cost: {
+    input: number
+    output: number
+  }
+  latency: number
+  tokens: {
+    input: number
+    output: number
+  }
+}

package/src/zai.ts

@@ -1,3 +1,4 @@
+import { Client } from '@botpress/client'
 import { BotpressClientLike, Cognitive, Model } from '@botpress/cognitive'
 
 import { type TextTokenizer, getWasmTokenizer } from '@bpinternal/thicktoken'
@@ -9,8 +10,13 @@ import { MemoryAdapter } from './adapters/memory'
 
 type ModelId = Required<Parameters<Cognitive['generateContent']>[0]['model']>
 
-type ActiveLearning = (typeof ActiveLearning)['_input']
-const ActiveLearning = z.object({
+type ActiveLearning = {
+  enable: boolean
+  tableName: string
+  taskId: string
+}
+
+const _ActiveLearning = z.object({
   enable: z.boolean().describe('Whether to enable active learning').default(false),
   tableName: z
     .string()
@@ -30,8 +36,15 @@ const ActiveLearning = z.object({
     .default('default'),
 })
 
-type ZaiConfig = (typeof ZaiConfig)['_input']
-const ZaiConfig = z.object({
+type ZaiConfig = {
+  client: BotpressClientLike | Cognitive
+  userId?: string
+  modelId?: ModelId | string
+  activeLearning?: ActiveLearning
+  namespace?: string
+}
+
+const _ZaiConfig = z.object({
   client: z.custom<BotpressClientLike | Cognitive>(),
   userId: z.string().describe('The ID of the user consuming the API').optional(),
   modelId: z
@@ -53,7 +66,7 @@ const ZaiConfig = z.object({
     )
     .describe('The ID of the model you want to use')
     .default('best' satisfies ModelId),
-  activeLearning: ActiveLearning.default({ enable: false }),
+  activeLearning: _ActiveLearning.default({ enable: false }),
   namespace: z
     .string()
     .regex(
@@ -79,7 +92,7 @@ export class Zai {
 
   public constructor(config: ZaiConfig) {
     this._originalConfig = config
-    const parsed = ZaiConfig.parse(config)
+    const parsed = _ZaiConfig.parse(config)
 
     this.client = Cognitive.isCognitiveClient(parsed.client)
       ? (parsed.client as unknown as Cognitive)
@@ -88,10 +101,13 @@ export class Zai {
     this.namespace = parsed.namespace
     this._userId = parsed.userId
     this.Model = parsed.modelId as ModelId
-    this.activeLearning = parsed.activeLearning
+    this.activeLearning = parsed.activeLearning as ActiveLearning
 
     this.adapter = parsed.activeLearning?.enable
-      ? new TableAdapter({ client: this.client.client, tableName: parsed.activeLearning.tableName })
+      ? new TableAdapter({
+          client: this.client.client as unknown as Client,
+          tableName: parsed.activeLearning.tableName,
+        })
       : new MemoryAdapter([])
   }