@botpress/zai 2.0.15 → 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>>;
     }
 }
 
@@ -172,16 +265,17 @@ type Options$1 = {
     instructions?: string;
     /** The maximum number of tokens per chunk */
     chunkLength?: number;
+    /** Whether to strictly follow the schema or not */
+    strict?: boolean;
 };
 type __Z<T extends any = any> = {
     _output: T;
 };
 type OfType<O, T extends __Z = __Z<O>> = T extends __Z<O> ? T : never;
-type AnyObjectOrArray = Record<string, unknown> | Array<unknown>;
 declare module '@botpress/zai' {
     interface Zai {
         /** Extracts one or many elements from an arbitrary input */
-        extract<S extends OfType<AnyObjectOrArray>>(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']>;
     }
 }
 
@@ -212,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);
+};