@botpress/zai 1.0.1 → 1.2.0

package/src/operations/text.ts

@@ -0,0 +1,65 @@
+// eslint-disable consistent-type-definitions
+import { z } from '@bpinternal/zui'
+
+import { clamp } from 'lodash-es'
+import { Zai } from '../zai'
+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'),
+})
+
+declare module '@botpress/zai' {
+  interface Zai {
+    /** Generates a text of the desired length according to the prompt */
+    text(prompt: string, options?: Options): Promise<string>
+  }
+}
+
+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.ModelDetails.input.maxTokens - PROMPT_INPUT_BUFFER, 100))
+
+  if (options.length) {
+    options.length = Math.min(this.ModelDetails.output.maxTokens - PROMPT_OUTPUT_BUFFER, options.length)
+  }
+
+  const instructions: string[] = []
+  let chart = ''
+
+  if (options.length) {
+    const length = clamp(options.length * 0.75, 5, options.length)
+    instructions.push(`IMPORTANT: Length constraint: ${length} tokens/words`)
+    instructions.push(`The text must be standalone and complete in less than ${length} tokens/words`)
+  }
+
+  if (options.length && options.length <= 500) {
+    chart = `
+| Tokens      | Text Length (approximate)            |
+|-------------|--------------------------------------|
+| < 5 tokens  | 1-3 words                            |
+| 5-10 tokens | 3-6 words                            |
+| 10-20 tokens| 6-15 words                           |
+| 20-50 tokens| A short sentence (15-30 words)       |
+| 50-100 tokens| A medium sentence (30-70 words)     |
+| 100-200 tokens| A short paragraph (70-150 words)   |
+| 200-300 tokens| A medium paragraph (150-200 words) |
+| 300-500 tokens| A long paragraph (200-300 words)   |`.trim()
+  }
+
+  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')}
+${chart}
+`.trim(),
+    temperature: 0.7,
+    messages: [{ type: 'text', content: prompt, role: 'user' }],
+    maxTokens: options.length,
+  })
+  return output?.choices?.[0]?.content! as string
+}

package/src/utils.ts

@@ -0,0 +1,52 @@
+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>'
+}
+
+export function fastHash(str: string): string {
+  let hash = 0
+  for (let i = 0; i < str.length; i++) {
+    hash = (hash << 5) - hash + str.charCodeAt(i)
+    hash |= 0 // Convert to 32bit integer
+  }
+  return (hash >>> 0).toString(16) // Convert to unsigned and then to hex
+}
+
+export const takeUntilTokens = <T>(arr: T[], tokens: number, count: (el: T) => number) => {
+  const result: T[] = []
+  let total = 0
+
+  for (const value of arr) {
+    const valueTokens = count(value)
+    if (total + valueTokens > tokens) {
+      break
+    }
+    total += valueTokens
+    result.push(value)
+  }
+
+  return result
+}
+
+export type GenerationMetadata = z.input<typeof GenerationMetadata>
+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/src/zai.ts

@@ -0,0 +1,147 @@
+import { BotpressClientLike, Cognitive, Model } from '@botpress/cognitive'
+
+import { type TextTokenizer, getWasmTokenizer } from '@bpinternal/thicktoken'
+import { z } from '@bpinternal/zui'
+
+import { Adapter } from './adapters/adapter'
+import { TableAdapter } from './adapters/botpress-table'
+import { MemoryAdapter } from './adapters/memory'
+
+type ModelId = Required<Parameters<Cognitive['generateContent']>[0]['model']>
+
+type ActiveLearning = z.input<typeof ActiveLearning>
+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$/,
+      'Namespace must be alphanumeric and contain only letters, numbers, underscores, hyphens and slashes'
+    )
+    .describe('The name of the table to store active learning tasks')
+    .default('ActiveLearningTable'),
+  taskId: z
+    .string()
+    .regex(
+      /^[A-Za-z0-9_/-]{1,100}$/,
+      'Namespace must be alphanumeric and contain only letters, numbers, underscores, hyphens and slashes'
+    )
+    .describe('The ID of the task')
+    .default('default'),
+})
+
+type ZaiConfig = z.input<typeof ZaiConfig>
+const ZaiConfig = z.object({
+  client: z.custom<BotpressClientLike | Cognitive>(),
+  userId: z.string().describe('The ID of the user consuming the API').optional(),
+  modelId: z
+    .custom<ModelId | string>(
+      (value) => {
+        if (typeof value !== 'string') {
+          return false
+        }
+
+        if (value !== 'best' && value !== 'fast' && !value.includes(':')) {
+          return false
+        }
+
+        return true
+      },
+      {
+        message: 'Invalid model ID',
+      }
+    )
+    .describe('The ID of the model you want to use')
+    .default('best' satisfies ModelId),
+  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'
+    )
+    .default('zai'),
+})
+
+export class Zai {
+  protected static tokenizer: TextTokenizer = null!
+  protected client: Cognitive
+
+  private _originalConfig: ZaiConfig
+
+  private _userId: string | undefined
+
+  protected Model: ModelId
+  protected ModelDetails: Model
+  protected namespace: string
+  protected adapter: Adapter
+  protected activeLearning: ActiveLearning
+
+  public constructor(config: ZaiConfig) {
+    this._originalConfig = config
+    const parsed = ZaiConfig.parse(config)
+
+    this.client = Cognitive.isCognitiveClient(parsed.client)
+      ? (parsed.client as unknown as Cognitive)
+      : new Cognitive({ client: parsed.client })
+
+    this.namespace = parsed.namespace
+    this._userId = parsed.userId
+    this.Model = parsed.modelId as ModelId
+    this.activeLearning = parsed.activeLearning
+
+    this.adapter = parsed.activeLearning?.enable
+      ? new TableAdapter({ client: this.client.client, tableName: parsed.activeLearning.tableName })
+      : new MemoryAdapter([])
+  }
+
+  /** @internal */
+  protected async callModel(
+    props: Parameters<Cognitive['generateContent']>[0]
+  ): ReturnType<Cognitive['generateContent']> {
+    return this.client.generateContent({
+      ...props,
+      model: this.Model,
+      userId: this._userId,
+    })
+  }
+
+  protected async getTokenizer() {
+    Zai.tokenizer ??= await (async () => {
+      while (!getWasmTokenizer) {
+        // there's an issue with wasm, it doesn't load immediately
+        await new Promise((resolve) => setTimeout(resolve, 25))
+      }
+      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
+    }
+
+    return `${this.namespace}/${this.activeLearning.taskId}`.replace(/\/+/g, '/')
+  }
+
+  public with(options: Partial<ZaiConfig>): Zai {
+    return new Zai({
+      ...this._originalConfig,
+      ...options,
+    })
+  }
+
+  public learn(taskId: string) {
+    return new Zai({
+      ...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",
-    "allowJs": true,
-    "skipLibCheck": true,
-    "esModuleInterop": true,
-    "allowSyntheticDefaultImports": true,
-    "forceConsistentCasingInFileNames": true,
-    "disableReferencedProjectLoad": true,
-    "resolveJsonModule": true,
-    "isolatedModules": true,
+    "outDir": "dist",
     "strict": false,
-    "noUnusedLocals": true,
-    "noUnusedParameters": true,
-    "noUncheckedIndexedAccess": true,
-    "lib": ["dom", "ESNext", "dom.iterable"],
-    "declaration": false,
-    "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/**/*"]
 }