@botpress/zai 2.0.15 → 2.1.0

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.15",
+  "version": "2.1.0",
   "main": "./dist/index.js",
   "types": "./dist/index.d.ts",
   "exports": {
@@ -11,7 +11,7 @@
   },
   "scripts": {
     "check:type": "tsc --noEmit",
-    "build": "bp add -y && pnpm run build:types && pnpm run build:neutral",
+    "build": "bp add -y && pnpm run build:types && pnpm run build:neutral && size-limit",
     "build:neutral": "ts-node -T ./build.ts",
     "build:types": "tsup",
     "watch": "tsup --watch",
@@ -19,11 +19,17 @@
     "test:e2e:update": "vitest -u run --config vitest.config.ts",
     "test:e2e:watch": "vitest --config vitest.config.ts"
   },
+  "size-limit": [
+    {
+      "limit": "50 kB",
+      "path": "dist/**/*.js"
+    }
+  ],
   "keywords": [],
   "author": "",
   "license": "ISC",
   "dependencies": {
-    "@botpress/cognitive": "0.1.28",
+    "@botpress/cognitive": "0.1.29",
     "json5": "^2.2.3",
     "jsonrepair": "^3.10.0",
     "lodash-es": "^4.17.21"
@@ -32,12 +38,14 @@
     "@botpress/client": "workspace:^",
     "@botpress/common": "workspace:*",
     "@botpress/vai": "workspace:*",
+    "@size-limit/file": "^11.1.6",
     "@types/lodash-es": "^4.17.12",
     "diff": "^8.0.1",
     "dotenv": "^16.4.4",
     "esbuild": "^0.16.12",
     "glob": "^9.3.4",
     "lodash": "^4.17.21",
+    "size-limit": "^11.1.6",
     "tsup": "^8.0.2"
   },
   "peerDependencies": {

package/src/context.ts

@@ -0,0 +1,197 @@
+import { Cognitive, Model, GenerateContentInput, GenerateContentOutput } from '@botpress/cognitive'
+import { Adapter } from './adapters/adapter'
+import { EventEmitter } from './emitter'
+
+type Meta = Awaited<ReturnType<Cognitive['generateContent']>>['meta']
+
+type GenerateContentProps<T> = Omit<GenerateContentInput, 'model' | 'signal'> & {
+  maxRetries?: number
+  transform?: (text: string | undefined, output: GenerateContentOutput) => T
+}
+
+export type ZaiContextProps = {
+  client: Cognitive
+  taskType: string
+  taskId: string
+  modelId: string
+  adapter?: Adapter
+  source?: GenerateContentInput['meta']
+}
+
+export 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
+}
+
+export class ZaiContext {
+  private _startedAt = Date.now()
+
+  private _inputCost = 0
+  private _outputCost = 0
+  private _inputTokens = 0
+  private _outputTokens = 0
+  private _totalCachedResponses = 0
+
+  private _totalRequests = 0
+  private _totalErrors = 0
+  private _totalResponses = 0
+
+  public taskId: string
+  public taskType: string
+  public modelId: GenerateContentInput['model']
+  public adapter?: Adapter
+  public source?: GenerateContentInput['meta']
+
+  private _eventEmitter: EventEmitter<ContextEvents>
+
+  public controller: AbortController = new AbortController()
+  private _client: Cognitive
+
+  public constructor(props: ZaiContextProps) {
+    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<ContextEvents>()
+
+    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)
+    })
+  }
+
+  public async getModel(): Promise<Model> {
+    return this._client.getModelDetails(this.modelId)
+  }
+
+  public on<K extends keyof ContextEvents>(type: K, listener: (event: ContextEvents[K]) => void) {
+    this._eventEmitter.on(type, listener)
+    return this
+  }
+
+  public clear() {
+    this._eventEmitter.clear()
+  }
+
+  public async generateContent<Out = string>(
+    props: GenerateContentProps<Out>
+  ): Promise<{ meta: Meta; output: GenerateContentOutput; text: string | undefined; extracted: Out }> {
+    const maxRetries = Math.max(props.maxRetries ?? 3, 0)
+    const transform = props.transform
+    let lastError: Error | null = 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: Out
+
+        messages.push({
+          role: 'assistant',
+          content: str || '<Invalid output, no content provided>',
+        })
+
+        if (!transform) {
+          output = str as any
+        } else {
+          output = transform(str, response.output)
+        }
+
+        return { meta: response.meta, output: response.output, text: str, extracted: output }
+      } catch (error) {
+        lastError = error as Error
+
+        if (attempt === maxRetries) {
+          throw lastError
+        }
+
+        messages.push({
+          role: 'user',
+          content: `ERROR PARSING OUTPUT\n\n${lastError.message}.\n\nPlease return a valid response addressing the error above.`,
+        })
+      }
+    }
+
+    throw lastError
+  }
+
+  public get elapsedTime(): number {
+    return Date.now() - this._startedAt
+  }
+
+  public get usage(): 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/src/emitter.ts

@@ -0,0 +1,49 @@
+export class EventEmitter<E extends object> {
+  private _listeners: {
+    [K in keyof E]?: ((event: E[K]) => void)[]
+  } = {}
+
+  public emit<K extends keyof E>(type: K, event: E[K]) {
+    const listeners = this._listeners[type]
+    if (!listeners) {
+      return
+    }
+    for (const listener of listeners) {
+      listener(event)
+    }
+  }
+
+  public once<K extends keyof E>(type: K, listener: (event: E[K]) => void) {
+    const wrapped = (event: E[K]) => {
+      this.off(type, wrapped)
+      listener(event)
+    }
+    this.on(type, wrapped)
+  }
+
+  public on<K extends keyof E>(type: K, listener: (event: E[K]) => void) {
+    if (!this._listeners[type]) {
+      this._listeners[type] = []
+    }
+    this._listeners[type]!.push(listener)
+  }
+
+  public off<K extends keyof E>(type: K, listener: (event: E[K]) => void) {
+    const listeners = this._listeners[type]
+    if (!listeners) {
+      return
+    }
+    const index = listeners.indexOf(listener)
+    if (index !== -1) {
+      listeners.splice(index, 1)
+    }
+  }
+
+  public clear<K extends keyof E>(type?: K) {
+    if (type) {
+      delete this._listeners[type]
+    } else {
+      this._listeners = {}
+    }
+  }
+}

package/src/operations/check.ts

@@ -1,6 +1,9 @@
 // eslint-disable consistent-type-definitions
 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'
@@ -35,12 +38,15 @@ declare module '@botpress/zai' {
       input: unknown,
       condition: string,
       options?: Options
-    ): Promise<{
-      /** Whether the condition is true or not */
-      value: boolean
-      /** The explanation of the decision */
-      explanation: string
-    }>
+    ): Response<
+      {
+        /** Whether the condition is true or not */
+        value: boolean
+        /** The explanation of the decision */
+        explanation: string
+      },
+      boolean
+    >
   }
 }
 
@@ -48,13 +54,21 @@ const TRUE = '■TRUE■'
 const FALSE = '■FALSE■'
 const END = '■END■'
 
-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)
-
-  const taskId = this.taskId
+const check = async (
+  input: unknown,
+  condition: string,
+  options: Options,
+  ctx: ZaiContext
+): Promise<{
+  value: boolean
+  explanation: string
+}> => {
+  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 = {
@@ -78,13 +92,14 @@ Zai.prototype.check = async function (this: Zai, input: unknown, condition: stri
     })
   )
 
-  const examples = taskId
-    ? await this.adapter.getExamples<string, boolean>({
-        input: inputAsString,
-        taskType,
-        taskId,
-      })
-    : []
+  const examples =
+    taskId && ctx.adapter
+      ? await ctx.adapter.getExamples<string, boolean>({
+          input: inputAsString,
+          taskType,
+          taskId,
+        })
+      : []
 
   const exactMatch = examples.find((x) => x.key === Key)
   if (exactMatch) {
@@ -163,7 +178,10 @@ ${END}
 `.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.
@@ -184,35 +202,36 @@ In your "Analysis", please refer to the Expert Examples # to justify your decisi
         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 finalAnswer: boolean
+      const explanation = text
+        .replace(TRUE, '')
+        .replace(FALSE, '')
+        .replace(END, '')
+        .replace('Final Answer:', '')
+        .replace('Analysis:', '')
+        .trim()
+
+      if (hasTrue && hasFalse) {
+        // If both TRUE and FALSE are present, we need to check which one was answered last
+        finalAnswer = text.lastIndexOf(TRUE) > text.lastIndexOf(FALSE)
+      } else {
+        finalAnswer = hasTrue
+      }
+
+      return { finalAnswer, explanation: explanation.trim() }
+    },
   })
 
-  const answer = output.choices[0]?.content as string
-
-  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: boolean
-  const explanation = answer
-    .replace(TRUE, '')
-    .replace(FALSE, '')
-    .replace(END, '')
-    .replace('Final Answer:', '')
-    .replace('Analysis:', '')
-    .trim()
-
-  if (hasTrue && hasFalse) {
-    // If both TRUE and FALSE are present, we need to check which one was answered last
-    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,
@@ -224,7 +243,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,
@@ -240,3 +259,34 @@ In your "Analysis", please refer to the Expert Examples # to justify your decisi
     explanation: explanation.trim(),
   }
 }
+
+Zai.prototype.check = function (
+  this: Zai,
+  input: unknown,
+  condition: string,
+  _options: Options | undefined
+): Response<
+  {
+    value: boolean
+    explanation: string
+  },
+  boolean
+> {
+  const options = _Options.parse(_options ?? {}) as Options
+
+  const context = new ZaiContext({
+    client: this.client,
+    modelId: this.Model,
+    taskId: this.taskId,
+    taskType: 'zai.check',
+    adapter: this.adapter,
+  })
+
+  return new Response<
+    {
+      value: boolean
+      explanation: string
+    },
+    boolean
+  >(context, check(input, condition, options, context), (result) => result.value)
+}