@botpress/zai 1.1.0 → 2.0.0

package/src/operations/rewrite.ts

@@ -1,3 +1,4 @@
+// eslint-disable consistent-type-definitions
 import { z } from '@bpinternal/zui'
 
 import { fastHash, stringify, takeUntilTokens } from '../utils'
@@ -7,13 +8,13 @@ import { PROMPT_INPUT_BUFFER } from './constants'
 type Example = z.input<typeof Example> & { instructions?: string }
 const Example = z.object({
   input: z.string(),
-  output: z.string()
+  output: z.string(),
 })
 
 export type Options = z.input<typeof Options>
 const Options = z.object({
   examples: z.array(Example).default([]),
-  length: z.number().min(10).max(16_000).optional().describe('The maximum number of tokens to generate')
+  length: z.number().min(10).max(16_000).optional().describe('The maximum number of tokens to generate'),
 })
 
 declare module '@botpress/zai' {
@@ -29,18 +30,19 @@ const END = '■END■'
 Zai.prototype.rewrite = async function (this: Zai, 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)`
     )
   }
 
@@ -69,27 +71,27 @@ ${END}
       taskId,
       taskType,
       input: original,
-      prompt
+      prompt,
     })
   )
 
   const formatExample = ({ input, output, instructions }: Example) => {
     return [
       { type: 'text' as const, role: 'user' as const, content: format(input, instructions || prompt) },
-      { type: 'text' as const, role: 'assistant' as const, content: `${START}${output}${END}` }
+      { type: 'text' as const, role: 'assistant' as const, content: `${START}${output}${END}` },
     ]
   }
 
   const defaultExamples: Example[] = [
     { input: 'Hello, how are you?', output: 'Bonjour, comment ça va?', instructions: 'translate to French' },
-    { input: '1\n2\n3', output: '3\n2\n1', instructions: 'reverse the order' }
+    { input: '1\n2\n3', output: '3\n2\n1', instructions: 'reverse the order' },
   ]
 
   const tableExamples = taskId
     ? await this.adapter.getExamples<string, string>({
         input: original,
         taskId,
-        taskType
+        taskType,
       })
     : []
 
@@ -100,10 +102,10 @@ ${END}
 
   const savedExamples: Example[] = [
     ...tableExamples.map((x) => ({ input: x.input as string, output: x.output as string })),
-    ...options.examples
+    ...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,
@@ -112,14 +114,14 @@ ${END}
     .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) => `• ${x}`).join('\n')}
 `.trim(),
     messages: [...examples, { type: 'text', content: format(original, prompt), role: 'user' }],
     maxTokens: options.length,
-    stopSequences: [END]
+    stopSequences: [END],
   })
 
   let result = output.choices[0]?.content as string
@@ -135,12 +137,23 @@ ${instructions.map((x) => `• ${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,
       taskType,
-      taskId
+      taskId,
     })
   }
 

package/src/operations/summarize.ts

@@ -1,3 +1,4 @@
+// eslint-disable consistent-type-definitions
 import { z } from '@bpinternal/zui'
 
 import { chunk } from 'lodash-es'
@@ -27,10 +28,10 @@ const Options = z.object({
   sliding: z
     .object({
       window: z.number().min(10).max(100_000),
-      overlap: z.number().min(0).max(100_000)
+      overlap: z.number().min(0).max(100_000),
     })
     .describe('Sliding window options')
-    .default({ window: 50_000, overlap: 250 })
+    .default({ window: 50_000, overlap: 250 }),
 })
 
 declare module '@botpress/zai' {
@@ -46,20 +47,21 @@ const END = '■END■'
 Zai.prototype.summarize = async function (this: Zai, original, _options) {
   const options = Options.parse(_options ?? {})
   const tokenizer = await this.getTokenizer()
+  await this.fetchModelDetails()
 
-  const INPUT_COMPONENT_SIZE = Math.max(100, (this.Model.input.maxTokens - PROMPT_INPUT_BUFFER) / 4)
+  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})`
     )
   }
 
   // Ensure the sliding window is not bigger than the model input size
-  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)
 
   // Ensure the overlap is not bigger than the window
   // Most extreme case possible (all 3 same size)
@@ -117,7 +119,7 @@ ${newText}
       'Summarize the text and make sure that the main points are included.',
       'Ignore any unnecessary details and focus on the main points.',
       'Use short and concise sentences to increase readability and information density.',
-      'When looking at the new information, focus on: ' + options.prompt
+      'When looking at the new information, focus on: ' + options.prompt,
     ]
 
     if (isFirst) {
@@ -157,7 +159,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).
@@ -171,7 +173,7 @@ ${options.format}
 `.trim(),
       messages: [{ type: 'text', content: format(currentSummary, slice), role: 'user' }],
       maxTokens: generationLength,
-      stopSequences: [END]
+      stopSequences: [END],
     })
 
     let result = output?.choices[0]?.content as string

package/src/operations/text.ts

@@ -1,3 +1,4 @@
+// eslint-disable consistent-type-definitions
 import { z } from '@bpinternal/zui'
 
 import { clamp } from 'lodash-es'
@@ -6,7 +7,7 @@ import { PROMPT_INPUT_BUFFER, PROMPT_OUTPUT_BUFFER } from './constants'
 
 export type Options = z.input<typeof Options>
 const Options = z.object({
-  length: z.number().min(1).max(100_000).optional().describe('The maximum number of tokens to generate')
+  length: z.number().min(1).max(100_000).optional().describe('The maximum number of tokens to generate'),
 })
 
 declare module '@botpress/zai' {
@@ -19,11 +20,12 @@ declare module '@botpress/zai' {
 Zai.prototype.text = async function (this: Zai, prompt, _options) {
   const options = Options.parse(_options ?? {})
   const tokenizer = await this.getTokenizer()
+  await this.fetchModelDetails()
 
-  prompt = tokenizer.truncate(prompt, Math.max(this.Model.input.maxTokens - PROMPT_INPUT_BUFFER, 100))
+  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: string[] = []
@@ -49,7 +51,7 @@ Zai.prototype.text = async function (this: Zai, prompt, _options) {
 | 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')}
@@ -57,7 +59,7 @@ ${chart}
 `.trim(),
     temperature: 0.7,
     messages: [{ type: 'text', content: prompt, role: 'user' }],
-    maxTokens: options.length
+    maxTokens: options.length,
   })
   return output?.choices?.[0]?.content! as string
 }

package/src/utils.ts

@@ -1,22 +1,13 @@
-import type { Client } from '@botpress/client'
 import { z } from '@bpinternal/zui'
 
 export const stringify = (input: unknown, beautify = true) => {
   return typeof input === 'string' && !!input.length
     ? input
     : input
-    ? JSON.stringify(input, beautify ? null : undefined, beautify ? 2 : undefined)
-    : '<input is null, false, undefined or empty>'
+      ? JSON.stringify(input, beautify ? null : undefined, beautify ? 2 : undefined)
+      : '<input is null, false, undefined or empty>'
 }
 
-export const BotpressClient = z.custom<Client | any>(
-  (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: string): string {
   let hash = 0
   for (let i = 0; i < str.length; i++) {
@@ -48,14 +39,14 @@ export const GenerationMetadata = z.object({
   cost: z
     .object({
       input: z.number(),
-      output: 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()
+      output: z.number(),
     })
-    .describe('Number of tokens used')
+    .describe('Number of tokens used'),
 })

package/src/zai.ts

@@ -1,16 +1,13 @@
-import { Client } from '@botpress/client'
+import { BotpressClientLike, Cognitive, Model } from '@botpress/cognitive'
+
+import { type TextTokenizer, getWasmTokenizer } from '@bpinternal/thicktoken'
 import { z } from '@bpinternal/zui'
-import { type TextTokenizer, getWasmTokenizer } from '@botpress/wasm'
 
 import { Adapter } from './adapters/adapter'
 import { TableAdapter } from './adapters/botpress-table'
 import { MemoryAdapter } from './adapters/memory'
-import { Models } from './models'
-import { llm } from './sdk-interfaces/llm/generateContent'
-
-import { BotpressClient, GenerationMetadata } from './utils'
 
-type ModelId = (typeof Models)[number]['id']
+type ModelId = Required<Parameters<Cognitive['generateContent']>[0]['model']>
 
 type ActiveLearning = z.input<typeof ActiveLearning>
 const ActiveLearning = z.object({
@@ -30,29 +27,32 @@ 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')
+    .default('default'),
 })
 
 type ZaiConfig = z.input<typeof ZaiConfig>
 const ZaiConfig = z.object({
-  client: BotpressClient,
+  client: z.custom<BotpressClientLike | Cognitive>(),
   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<ModelId | string>(
       (value) => {
-        if (typeof value !== 'string' || !value.includes('__')) {
+        if (typeof value !== 'string') {
+          return false
+        }
+
+        if (value !== 'best' && value !== 'fast' && !value.includes(':')) {
           return false
         }
 
         return true
       },
       {
-        message: 'Invalid model ID'
+        message: 'Invalid model ID',
       }
     )
     .describe('The ID of the model you want to use')
-    .default('openai__gpt-4o-mini-2024-07-18' satisfies ModelId),
+    .default('best' satisfies ModelId),
   activeLearning: ActiveLearning.default({ enable: false }),
   namespace: z
     .string()
@@ -60,102 +60,50 @@ const ZaiConfig = z.object({
       /^[A-Za-z0-9_/-]{1,100}$/,
       'Namespace must be alphanumeric and contain only letters, numbers, underscores, hyphens and slashes'
     )
-    .default('zai')
+    .default('zai'),
 })
 
 export class Zai {
   protected static tokenizer: TextTokenizer = null!
-  protected client: Client
+  protected client: Cognitive
 
-  private originalConfig: ZaiConfig
+  private _originalConfig: ZaiConfig
 
-  private userId: string | undefined
-  private integration: string
-  private model: string
-  private retry: { maxRetries: number }
+  private _userId: string | undefined
 
-  protected Model: (typeof Models)[number]
+  protected Model: ModelId
+  protected ModelDetails: Model
   protected namespace: string
   protected adapter: Adapter
   protected activeLearning: ActiveLearning
 
-  constructor(config: ZaiConfig) {
-    this.originalConfig = config
+  public constructor(config: ZaiConfig) {
+    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.client = Cognitive.isCognitiveClient(parsed.client)
+      ? (parsed.client as unknown as Cognitive)
+      : new Cognitive({ client: parsed.client })
 
-    this.integration = integration!
-    this.model = modelId!
     this.namespace = parsed.namespace
-    this.userId = parsed.userId
-    this.retry = parsed.retry as { maxRetries: number }
-    this.Model = Models.find((m) => m.id === parsed.modelId)!
+    this._userId = parsed.userId
+    this.Model = parsed.modelId as ModelId
     this.activeLearning = parsed.activeLearning
 
     this.adapter = parsed.activeLearning?.enable
-      ? new TableAdapter({ client: this.client, tableName: parsed.activeLearning.tableName })
+      ? new TableAdapter({ client: this.client.client, tableName: parsed.activeLearning.tableName })
       : new MemoryAdapter([])
   }
 
   /** @internal */
   protected async callModel(
-    props: Partial<llm.generateContent.Input>
-  ): Promise<llm.generateContent.Output & { metadata: GenerationMetadata }> {
-    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, 1000))
-        } else {
-          throw new Error('Failed to call model after multiple retries')
-        }
-      }
-    }
-
-    throw new Error('Failed to call model after multiple retries')
-  }
-
-  /** @internal */
-  private async _callModel(
-    props: Partial<llm.generateContent.Input>
-  ): Promise<llm.generateContent.Output & { metadata: GenerationMetadata }> {
-    let retries = this.retry.maxRetries
-    do {
-      const start = Date.now()
-      const input: llm.generateContent.Input = {
-        messages: [],
-        temperature: 0.0,
-        topP: 1,
-        model: { id: this.model },
-        userId: this.userId,
-        ...props
-      }
-
-      const { output } = (await this.client.callAction({
-        type: `${this.integration}:generateContent`,
-        input
-      })) as unknown as { output: llm.generateContent.Output }
-
-      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)
+    props: Parameters<Cognitive['generateContent']>[0]
+  ): ReturnType<Cognitive['generateContent']> {
+    return this.client.generateContent({
+      ...props,
+      model: this.Model,
+      userId: this._userId,
+    })
   }
 
   protected async getTokenizer() {
@@ -164,11 +112,17 @@ export class Zai {
         // there's an issue with wasm, it doesn't load immediately
         await new Promise((resolve) => setTimeout(resolve, 25))
       }
-      return getWasmTokenizer()
+      return getWasmTokenizer() as TextTokenizer
     })()
     return Zai.tokenizer
   }
 
+  protected async fetchModelDetails(): Promise<void> {
+    if (!this.ModelDetails) {
+      this.ModelDetails = await this.client.getModelDetails(this.Model)
+    }
+  }
+
   protected get taskId() {
     if (!this.activeLearning.enable) {
       return undefined
@@ -179,15 +133,15 @@ export class Zai {
 
   public with(options: Partial<ZaiConfig>): Zai {
     return new Zai({
-      ...this.originalConfig,
-      ...options
+      ...this._originalConfig,
+      ...options,
     })
   }
 
   public learn(taskId: string) {
     return new Zai({
-      ...this.originalConfig,
-      activeLearning: { ...this.activeLearning, taskId, enable: true }
+      ...this._originalConfig,
+      activeLearning: { ...this.activeLearning, taskId, enable: true },
     })
   }
 }

package/tsconfig.json

@@ -1,32 +1,12 @@
 {
+  "extends": "../../tsconfig.json",
   "compilerOptions": {
-    "target": "ESNext",
-    "module": "ESNext",
-    "moduleResolution": "bundler",
     "outDir": "dist",
-    "allowJs": true,
-    "skipLibCheck": true,
-    "esModuleInterop": true,
-    "allowSyntheticDefaultImports": true,
-    "forceConsistentCasingInFileNames": true,
-    "disableReferencedProjectLoad": true,
-    "resolveJsonModule": true,
-    "isolatedModules": true,
     "strict": false,
-    "noUnusedLocals": true,
-    "noUnusedParameters": true,
-    "noUncheckedIndexedAccess": true,
-    "lib": ["dom", "ESNext", "dom.iterable"],
-    "declaration": true,
-    "noEmit": false,
     "paths": {
       "@botpress/zai": ["./src/zai.ts"]
     }
   },
   "exclude": ["node_modules", "dist"],
-  "include": ["src/**/*", "vitest.d.ts"],
-  "ts-node": {
-    "esm": true,
-    "require": ["dotenv/config", "./ensure-env.cjs"]
-  }
+  "include": ["src/**/*", "vitest.d.ts", "e2e/**/*"]
 }