@botpress/zai 1.0.1 → 1.1.0

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (47) hide show
  1. package/dist/adapters/adapter.js +2 -0
  2. package/dist/adapters/botpress-table.js +168 -0
  3. package/dist/adapters/memory.js +12 -0
  4. package/dist/index.d.ts +99 -98
  5. package/dist/index.js +9 -1873
  6. package/dist/models.js +387 -0
  7. package/dist/operations/check.js +141 -0
  8. package/dist/operations/constants.js +2 -0
  9. package/dist/operations/errors.js +15 -0
  10. package/dist/operations/extract.js +212 -0
  11. package/dist/operations/filter.js +179 -0
  12. package/dist/operations/label.js +237 -0
  13. package/dist/operations/rewrite.js +111 -0
  14. package/dist/operations/summarize.js +132 -0
  15. package/dist/operations/text.js +46 -0
  16. package/dist/utils.js +43 -0
  17. package/dist/zai.js +140 -0
  18. package/package.json +21 -19
  19. package/src/adapters/adapter.ts +35 -0
  20. package/src/adapters/botpress-table.ts +210 -0
  21. package/src/adapters/memory.ts +13 -0
  22. package/src/index.ts +11 -0
  23. package/src/models.ts +394 -0
  24. package/src/operations/__tests/botpress_docs.txt +26040 -0
  25. package/src/operations/__tests/cache.jsonl +101 -0
  26. package/src/operations/__tests/index.ts +87 -0
  27. package/src/operations/check.ts +187 -0
  28. package/src/operations/constants.ts +2 -0
  29. package/src/operations/errors.ts +9 -0
  30. package/src/operations/extract.ts +291 -0
  31. package/src/operations/filter.ts +231 -0
  32. package/src/operations/label.ts +332 -0
  33. package/src/operations/rewrite.ts +148 -0
  34. package/src/operations/summarize.ts +193 -0
  35. package/src/operations/text.ts +63 -0
  36. package/src/sdk-interfaces/llm/generateContent.ts +127 -0
  37. package/src/sdk-interfaces/llm/listLanguageModels.ts +19 -0
  38. package/src/utils.ts +61 -0
  39. package/src/zai.ts +193 -0
  40. package/tsconfig.json +2 -2
  41. package/dist/index.cjs +0 -1903
  42. package/dist/index.cjs.map +0 -1
  43. package/dist/index.d.cts +0 -916
  44. package/dist/index.js.map +0 -1
  45. package/tsup.config.ts +0 -16
  46. package/vitest.config.ts +0 -9
  47. package/vitest.setup.ts +0 -24
@@ -0,0 +1,193 @@
1
+ import { z } from '@bpinternal/zui'
2
+
3
+ import { chunk } from 'lodash-es'
4
+ import { Zai } from '../zai'
5
+ import { PROMPT_INPUT_BUFFER, PROMPT_OUTPUT_BUFFER } from './constants'
6
+
7
+ export type Options = z.input<typeof Options>
8
+ const Options = z.object({
9
+ prompt: z
10
+ .string()
11
+ .describe('What should the text be summarized to?')
12
+ .default('New information, concepts and ideas that are deemed important'),
13
+ format: z
14
+ .string()
15
+ .describe('How to format the example text')
16
+ .default(
17
+ 'A normal text with multiple sentences and paragraphs. Use markdown to format the text into sections. Use headings, lists, and other markdown features to make the text more readable. Do not include links, images, or other non-text elements.'
18
+ ),
19
+ length: z.number().min(10).max(100_000).describe('The length of the summary in tokens').default(250),
20
+ intermediateFactor: z
21
+ .number()
22
+ .min(1)
23
+ .max(10)
24
+ .describe('How many times longer (than final length) are the intermediate summaries generated')
25
+ .default(4),
26
+ maxIterations: z.number().min(1).default(100),
27
+ sliding: z
28
+ .object({
29
+ window: z.number().min(10).max(100_000),
30
+ overlap: z.number().min(0).max(100_000)
31
+ })
32
+ .describe('Sliding window options')
33
+ .default({ window: 50_000, overlap: 250 })
34
+ })
35
+
36
+ declare module '@botpress/zai' {
37
+ interface Zai {
38
+ /** Summarizes a text of any length to a summary of the desired length */
39
+ summarize(original: string, options?: Options): Promise<string>
40
+ }
41
+ }
42
+
43
+ const START = '■START■'
44
+ const END = '■END■'
45
+
46
+ Zai.prototype.summarize = async function (this: Zai, original, _options) {
47
+ const options = Options.parse(_options ?? {})
48
+ const tokenizer = await this.getTokenizer()
49
+
50
+ const INPUT_COMPONENT_SIZE = Math.max(100, (this.Model.input.maxTokens - PROMPT_INPUT_BUFFER) / 4)
51
+ options.prompt = tokenizer.truncate(options.prompt, INPUT_COMPONENT_SIZE)
52
+ options.format = tokenizer.truncate(options.format, INPUT_COMPONENT_SIZE)
53
+
54
+ const maxOutputSize = this.Model.output.maxTokens - PROMPT_OUTPUT_BUFFER
55
+ if (options.length > maxOutputSize) {
56
+ throw new Error(
57
+ `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})`
58
+ )
59
+ }
60
+
61
+ // Ensure the sliding window is not bigger than the model input size
62
+ options.sliding.window = Math.min(options.sliding.window, this.Model.input.maxTokens - PROMPT_INPUT_BUFFER)
63
+
64
+ // Ensure the overlap is not bigger than the window
65
+ // Most extreme case possible (all 3 same size)
66
+ // |ooooooooooooooooooo|wwwwwwwwwwwwwww|ooooooooooooooooooo|
67
+ // |<---- overlap ---->|<-- window -->|<---- overlap ---->|
68
+ options.sliding.overlap = Math.min(options.sliding.overlap, options.sliding.window - 3 * options.sliding.overlap)
69
+
70
+ const format = (summary: string, newText: string) => {
71
+ return `
72
+ ${START}
73
+ ${summary.length ? summary : '<summary still empty>'}
74
+ ${END}
75
+
76
+ Please amend the summary between the ${START} and ${END} tags to accurately reflect the prompt and the additional text below.
77
+
78
+ <|start_new_information|>
79
+ ${newText}
80
+ <|new_information|>`.trim()
81
+ }
82
+
83
+ const tokens = tokenizer.split(original)
84
+ const parts = Math.ceil(tokens.length / (options.sliding.window - options.sliding.overlap))
85
+ let iteration = 0
86
+
87
+ // We split it recursively into smaller parts until we're at less than 4 window slides per part
88
+ // Then we use a merge strategy to combine the sub-chunks summaries
89
+ // This is basically a merge sort algorithm (but summary instead of sorting)
90
+ const N = 2 // This is the merge sort exponent
91
+ const useMergeSort = parts >= Math.pow(2, N)
92
+ const chunkSize = Math.ceil(tokens.length / (parts * N))
93
+
94
+ if (useMergeSort) {
95
+ const chunks = chunk(tokens, chunkSize).map((x) => x.join(''))
96
+ const allSummaries = await Promise.all(chunks.map((chunk) => this.summarize(chunk, options)))
97
+ return this.summarize(allSummaries.join('\n\n============\n\n'), options)
98
+ }
99
+
100
+ const summaries: string[] = []
101
+ let currentSummary = ''
102
+
103
+ for (let i = 0; i < tokens.length; i += options.sliding.window) {
104
+ const from = Math.max(0, i - options.sliding.overlap)
105
+ const to = Math.min(tokens.length, i + options.sliding.window + options.sliding.overlap)
106
+ const isFirst = i === 0
107
+ const isLast = to >= tokens.length
108
+
109
+ const slice = tokens.slice(from, to).join('')
110
+
111
+ if (iteration++ >= options.maxIterations) {
112
+ break
113
+ }
114
+
115
+ const instructions: string[] = [
116
+ `At each step, you will receive a part of the text to summarize. Make sure to reply with the new summary in the tags ${START} and ${END}.`,
117
+ 'Summarize the text and make sure that the main points are included.',
118
+ 'Ignore any unnecessary details and focus on the main points.',
119
+ 'Use short and concise sentences to increase readability and information density.',
120
+ 'When looking at the new information, focus on: ' + options.prompt
121
+ ]
122
+
123
+ if (isFirst) {
124
+ instructions.push(
125
+ 'The current summary is empty. You need to generate a summary that covers the main points of the text.'
126
+ )
127
+ }
128
+
129
+ let generationLength = options.length
130
+
131
+ if (!isLast) {
132
+ generationLength = Math.min(
133
+ tokenizer.count(currentSummary) + options.length * options.intermediateFactor,
134
+ maxOutputSize
135
+ )
136
+
137
+ instructions.push(
138
+ 'You need to amend the summary to include the new information. Make sure the summary is complete and covers all the main points.'
139
+ )
140
+
141
+ instructions.push(`The current summary is ${currentSummary.length} tokens long.`)
142
+ instructions.push(`You can amend the summary to be up to ${generationLength} tokens long.`)
143
+ }
144
+
145
+ if (isLast) {
146
+ instructions.push(
147
+ 'This is the last part you will have to summarize. Make sure the summary is complete and covers all the main points.'
148
+ )
149
+ instructions.push(
150
+ `The current summary is ${currentSummary.length} tokens long. You need to make sure it is ${options.length} tokens or less.`
151
+ )
152
+
153
+ if (currentSummary.length > options.length) {
154
+ instructions.push(
155
+ `The current summary is already too long, so you need to shorten it to ${options.length} tokens while also including the new information.`
156
+ )
157
+ }
158
+ }
159
+
160
+ const output = await this.callModel({
161
+ systemPrompt: `
162
+ You are summarizing a text. The text is split into ${parts} parts, and you are currently working on part ${iteration}.
163
+ 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).
164
+ The summary needs to cover the main points of the text and must be concise.
165
+
166
+ IMPORTANT INSTRUCTIONS:
167
+ ${instructions.map((x) => `- ${x.trim()}`).join('\n')}
168
+
169
+ FORMAT OF THE SUMMARY:
170
+ ${options.format}
171
+ `.trim(),
172
+ messages: [{ type: 'text', content: format(currentSummary, slice), role: 'user' }],
173
+ maxTokens: generationLength,
174
+ stopSequences: [END]
175
+ })
176
+
177
+ let result = output?.choices[0]?.content as string
178
+
179
+ if (result.includes(START)) {
180
+ result = result.slice(result.indexOf(START) + START.length)
181
+ }
182
+
183
+ if (result.includes('■')) {
184
+ // can happen if the model truncates the text before the entire END tag is written
185
+ result = result.slice(0, result.indexOf('■'))
186
+ }
187
+
188
+ summaries.push(result)
189
+ currentSummary = result
190
+ }
191
+
192
+ return currentSummary.trim()
193
+ }
@@ -0,0 +1,63 @@
1
+ import { z } from '@bpinternal/zui'
2
+
3
+ import { clamp } from 'lodash-es'
4
+ import { Zai } from '../zai'
5
+ import { PROMPT_INPUT_BUFFER, PROMPT_OUTPUT_BUFFER } from './constants'
6
+
7
+ export type Options = z.input<typeof Options>
8
+ const Options = z.object({
9
+ length: z.number().min(1).max(100_000).optional().describe('The maximum number of tokens to generate')
10
+ })
11
+
12
+ declare module '@botpress/zai' {
13
+ interface Zai {
14
+ /** Generates a text of the desired length according to the prompt */
15
+ text(prompt: string, options?: Options): Promise<string>
16
+ }
17
+ }
18
+
19
+ Zai.prototype.text = async function (this: Zai, prompt, _options) {
20
+ const options = Options.parse(_options ?? {})
21
+ const tokenizer = await this.getTokenizer()
22
+
23
+ prompt = tokenizer.truncate(prompt, Math.max(this.Model.input.maxTokens - PROMPT_INPUT_BUFFER, 100))
24
+
25
+ if (options.length) {
26
+ options.length = Math.min(this.Model.output.maxTokens - PROMPT_OUTPUT_BUFFER, options.length)
27
+ }
28
+
29
+ const instructions: string[] = []
30
+ let chart = ''
31
+
32
+ if (options.length) {
33
+ const length = clamp(options.length * 0.75, 5, options.length)
34
+ instructions.push(`IMPORTANT: Length constraint: ${length} tokens/words`)
35
+ instructions.push(`The text must be standalone and complete in less than ${length} tokens/words`)
36
+ }
37
+
38
+ if (options.length && options.length <= 500) {
39
+ chart = `
40
+ | Tokens | Text Length (approximate) |
41
+ |-------------|--------------------------------------|
42
+ | < 5 tokens | 1-3 words |
43
+ | 5-10 tokens | 3-6 words |
44
+ | 10-20 tokens| 6-15 words |
45
+ | 20-50 tokens| A short sentence (15-30 words) |
46
+ | 50-100 tokens| A medium sentence (30-70 words) |
47
+ | 100-200 tokens| A short paragraph (70-150 words) |
48
+ | 200-300 tokens| A medium paragraph (150-200 words) |
49
+ | 300-500 tokens| A long paragraph (200-300 words) |`.trim()
50
+ }
51
+
52
+ const output = await this.callModel({
53
+ systemPrompt: `
54
+ 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.
55
+ ${instructions.map((x) => `- ${x}`).join('\n')}
56
+ ${chart}
57
+ `.trim(),
58
+ temperature: 0.7,
59
+ messages: [{ type: 'text', content: prompt, role: 'user' }],
60
+ maxTokens: options.length
61
+ })
62
+ return output?.choices?.[0]?.content! as string
63
+ }
@@ -0,0 +1,127 @@
1
+
2
+ // This file is generated. Do not edit it manually.
3
+ // See 'scripts/update-models.ts'
4
+
5
+ /* eslint-disable */
6
+ /* tslint:disable */
7
+
8
+ export namespace llm {
9
+ export namespace generateContent {
10
+ export type Input = {
11
+ /** Model to use for content generation */
12
+ model?: { id: string }
13
+ ;
14
+ /** Optional system prompt to guide the model */
15
+ systemPrompt?: string
16
+ ;
17
+ /** Array of messages for the model to process */
18
+ messages: Array<{ role: 'user' | 'assistant'; type?: 'text' | 'tool_calls' | 'tool_result' | 'multipart';
19
+ /** Required if `type` is "tool_calls" */
20
+ toolCalls?: Array<{ id: string; type: 'function'; function: { name: string;
21
+ /** Some LLMs may generate invalid JSON for a tool call, so this will be `null` when it happens. */
22
+ arguments:
23
+ { [key: string]: any } | null
24
+ } }>
25
+ ;
26
+ /** Required if `type` is "tool_result" */
27
+ toolResultCallId?: string
28
+ ;
29
+ /** Required unless `type` is "tool_call". If `type` is "multipart", this field must be an array of content objects. If `type` is "tool_result" then this field should be the result of the tool call (a plain string or a JSON-encoded array or object). If `type` is "tool_call" then the `toolCalls` field should be used instead. */
30
+ content:
31
+
32
+ string | Array<{ type: 'text' | 'image';
33
+ /** Indicates the MIME type of the content. If not provided it will be detected from the content-type header of the provided URL. */
34
+ mimeType?: string
35
+ ;
36
+ /** Required if part type is "text" */
37
+ text?: string
38
+ ;
39
+ /** Required if part type is "image" */
40
+ url?: string
41
+ }> | null
42
+ }>
43
+ ;
44
+ /** Response format expected from the model. If "json_object" is chosen, you must instruct the model to generate JSON either via the system prompt or a user message. */
45
+ responseFormat?: 'text' | 'json_object'
46
+ ;
47
+ /** Maximum number of tokens allowed in the generated response */
48
+ maxTokens?: number
49
+ ;
50
+ /** Sampling temperature for the model. Higher values result in more random outputs. */
51
+ temperature?: /** Sampling temperature for the model. Higher values result in more random outputs. */ number
52
+ ;
53
+ /** Top-p sampling parameter. Limits sampling to the smallest set of tokens with a cumulative probability above the threshold. */
54
+ topP?: /** Top-p sampling parameter. Limits sampling to the smallest set of tokens with a cumulative probability above the threshold. */ number
55
+ ;
56
+ /** Sequences where the model should stop generating further tokens. */
57
+ stopSequences?: string[]
58
+ ; tools?: Array<{ type: 'function'; function: {
59
+ /** Function name */
60
+ name: string
61
+ ; description?: string;
62
+ /** JSON schema of the function arguments */
63
+ argumentsSchema?: { }
64
+ } }>; toolChoice?: { type?: 'auto' | 'specific' | 'any' | 'none' | '';
65
+ /** Required if `type` is "specific" */
66
+ functionName?: string
67
+ }; userId?: string;
68
+ /** Set to `true` to output debug information to the bot logs */
69
+ debug?: boolean
70
+ ; meta?: {
71
+ /** Source of the prompt, e.g. agent/:id/:version cards/ai-generate, cards/ai-task, nodes/autonomous, etc. */
72
+ promptSource?: string
73
+ ; promptCategory?: string;
74
+ /** Name of the integration that originally received the message that initiated this action */
75
+ integrationName?: string
76
+ } };;
77
+ export type Output = {
78
+ /** Response ID from LLM provider */
79
+ id: string
80
+ ;
81
+ /** LLM provider name */
82
+ provider: string
83
+ ;
84
+ /** Model name */
85
+ model: string
86
+ ; choices: Array<{ type?: 'text' | 'tool_calls' | 'tool_result' | 'multipart';
87
+ /** Required if `type` is "tool_calls" */
88
+ toolCalls?: Array<{ id: string; type: 'function'; function: { name: string;
89
+ /** Some LLMs may generate invalid JSON for a tool call, so this will be `null` when it happens. */
90
+ arguments:
91
+ { [key: string]: any } | null
92
+ } }>
93
+ ;
94
+ /** Required if `type` is "tool_result" */
95
+ toolResultCallId?: string
96
+ ;
97
+ /** Required unless `type` is "tool_call". If `type` is "multipart", this field must be an array of content objects. If `type` is "tool_result" then this field should be the result of the tool call (a plain string or a JSON-encoded array or object). If `type` is "tool_call" then the `toolCalls` field should be used instead. */
98
+ content:
99
+
100
+ string | Array<{ type: 'text' | 'image';
101
+ /** Indicates the MIME type of the content. If not provided it will be detected from the content-type header of the provided URL. */
102
+ mimeType?: string
103
+ ;
104
+ /** Required if part type is "text" */
105
+ text?: string
106
+ ;
107
+ /** Required if part type is "image" */
108
+ url?: string
109
+ }> | null
110
+ ; role: 'assistant'; index: number; stopReason: 'stop' | 'max_tokens' | 'tool_calls' | 'content_filter' | 'other' }>; usage: {
111
+ /** Number of input tokens used by the model */
112
+ inputTokens: number
113
+ ;
114
+ /** Cost of the input tokens received by the model, in U.S. dollars */
115
+ inputCost: number
116
+ ;
117
+ /** Number of output tokens used by the model */
118
+ outputTokens: number
119
+ ;
120
+ /** Cost of the output tokens generated by the model, in U.S. dollars */
121
+ outputCost: number
122
+ }; botpress: {
123
+ /** Total cost of the content generation, in U.S. dollars */
124
+ cost: number
125
+ } };;
126
+ }
127
+ }
@@ -0,0 +1,19 @@
1
+
2
+ // This file is generated. Do not edit it manually.
3
+ // See 'scripts/update-models.ts'
4
+
5
+ /* eslint-disable */
6
+ /* tslint:disable */
7
+
8
+ export namespace llm {
9
+ export namespace listLanguageModels {
10
+ export type Input = { };;
11
+ export type Output = { models: Array<{ id: string; name: string; description: string; tags: Array<'recommended' | 'deprecated' | 'general-purpose' | 'low-cost' | 'vision' | 'coding' | 'agents' | 'function-calling' | 'roleplay' | 'storytelling' | 'reasoning'>; input: { maxTokens: number;
12
+ /** Cost per 1 million tokens, in U.S. dollars */
13
+ costPer1MTokens: number
14
+ }; output: { maxTokens: number;
15
+ /** Cost per 1 million tokens, in U.S. dollars */
16
+ costPer1MTokens: number
17
+ } } & { id: string }> };;
18
+ }
19
+ }
package/src/utils.ts ADDED
@@ -0,0 +1,61 @@
1
+ import type { Client } from '@botpress/client'
2
+ import { z } from '@bpinternal/zui'
3
+
4
+ export const stringify = (input: unknown, beautify = true) => {
5
+ return typeof input === 'string' && !!input.length
6
+ ? input
7
+ : input
8
+ ? JSON.stringify(input, beautify ? null : undefined, beautify ? 2 : undefined)
9
+ : '<input is null, false, undefined or empty>'
10
+ }
11
+
12
+ export const BotpressClient = z.custom<Client | any>(
13
+ (value) =>
14
+ typeof value === 'object' && value !== null && 'callAction' in value && typeof value.callAction === 'function',
15
+ {
16
+ message: 'Invalid Botpress Client. Make sure to pass an instance of @botpress/client'
17
+ }
18
+ )
19
+
20
+ export function fastHash(str: string): string {
21
+ let hash = 0
22
+ for (let i = 0; i < str.length; i++) {
23
+ hash = (hash << 5) - hash + str.charCodeAt(i)
24
+ hash |= 0 // Convert to 32bit integer
25
+ }
26
+ return (hash >>> 0).toString(16) // Convert to unsigned and then to hex
27
+ }
28
+
29
+ export const takeUntilTokens = <T>(arr: T[], tokens: number, count: (el: T) => number) => {
30
+ const result: T[] = []
31
+ let total = 0
32
+
33
+ for (const value of arr) {
34
+ const valueTokens = count(value)
35
+ if (total + valueTokens > tokens) {
36
+ break
37
+ }
38
+ total += valueTokens
39
+ result.push(value)
40
+ }
41
+
42
+ return result
43
+ }
44
+
45
+ export type GenerationMetadata = z.input<typeof GenerationMetadata>
46
+ export const GenerationMetadata = z.object({
47
+ model: z.string(),
48
+ cost: z
49
+ .object({
50
+ input: z.number(),
51
+ output: z.number()
52
+ })
53
+ .describe('Cost in $USD'),
54
+ latency: z.number().describe('Latency in milliseconds'),
55
+ tokens: z
56
+ .object({
57
+ input: z.number(),
58
+ output: z.number()
59
+ })
60
+ .describe('Number of tokens used')
61
+ })
package/src/zai.ts ADDED
@@ -0,0 +1,193 @@
1
+ import { Client } from '@botpress/client'
2
+ import { z } from '@bpinternal/zui'
3
+ import { type TextTokenizer, getWasmTokenizer } from '@botpress/wasm'
4
+
5
+ import { Adapter } from './adapters/adapter'
6
+ import { TableAdapter } from './adapters/botpress-table'
7
+ import { MemoryAdapter } from './adapters/memory'
8
+ import { Models } from './models'
9
+ import { llm } from './sdk-interfaces/llm/generateContent'
10
+
11
+ import { BotpressClient, GenerationMetadata } from './utils'
12
+
13
+ type ModelId = (typeof Models)[number]['id']
14
+
15
+ type ActiveLearning = z.input<typeof ActiveLearning>
16
+ const ActiveLearning = z.object({
17
+ enable: z.boolean().describe('Whether to enable active learning').default(false),
18
+ tableName: z
19
+ .string()
20
+ .regex(
21
+ /^[A-Za-z0-9_/-]{1,100}Table$/,
22
+ 'Namespace must be alphanumeric and contain only letters, numbers, underscores, hyphens and slashes'
23
+ )
24
+ .describe('The name of the table to store active learning tasks')
25
+ .default('ActiveLearningTable'),
26
+ taskId: z
27
+ .string()
28
+ .regex(
29
+ /^[A-Za-z0-9_/-]{1,100}$/,
30
+ 'Namespace must be alphanumeric and contain only letters, numbers, underscores, hyphens and slashes'
31
+ )
32
+ .describe('The ID of the task')
33
+ .default('default')
34
+ })
35
+
36
+ type ZaiConfig = z.input<typeof ZaiConfig>
37
+ const ZaiConfig = z.object({
38
+ client: BotpressClient,
39
+ userId: z.string().describe('The ID of the user consuming the API').optional(),
40
+ retry: z.object({ maxRetries: z.number().min(0).max(100) }).default({ maxRetries: 3 }),
41
+ modelId: z
42
+ .custom<ModelId | string>(
43
+ (value) => {
44
+ if (typeof value !== 'string' || !value.includes('__')) {
45
+ return false
46
+ }
47
+
48
+ return true
49
+ },
50
+ {
51
+ message: 'Invalid model ID'
52
+ }
53
+ )
54
+ .describe('The ID of the model you want to use')
55
+ .default('openai__gpt-4o-mini-2024-07-18' satisfies ModelId),
56
+ activeLearning: ActiveLearning.default({ enable: false }),
57
+ namespace: z
58
+ .string()
59
+ .regex(
60
+ /^[A-Za-z0-9_/-]{1,100}$/,
61
+ 'Namespace must be alphanumeric and contain only letters, numbers, underscores, hyphens and slashes'
62
+ )
63
+ .default('zai')
64
+ })
65
+
66
+ export class Zai {
67
+ protected static tokenizer: TextTokenizer = null!
68
+ protected client: Client
69
+
70
+ private originalConfig: ZaiConfig
71
+
72
+ private userId: string | undefined
73
+ private integration: string
74
+ private model: string
75
+ private retry: { maxRetries: number }
76
+
77
+ protected Model: (typeof Models)[number]
78
+ protected namespace: string
79
+ protected adapter: Adapter
80
+ protected activeLearning: ActiveLearning
81
+
82
+ constructor(config: ZaiConfig) {
83
+ this.originalConfig = config
84
+ const parsed = ZaiConfig.parse(config)
85
+
86
+ this.client = parsed.client
87
+ const [integration, modelId] = parsed.modelId.split('__')
88
+
89
+ if (!integration?.length || !modelId?.length) {
90
+ throw new Error(`Invalid model ID: ${parsed.modelId}. Expected format: <integration>__<modelId>`)
91
+ }
92
+
93
+ this.integration = integration!
94
+ this.model = modelId!
95
+ this.namespace = parsed.namespace
96
+ this.userId = parsed.userId
97
+ this.retry = parsed.retry as { maxRetries: number }
98
+ this.Model = Models.find((m) => m.id === parsed.modelId)!
99
+ this.activeLearning = parsed.activeLearning
100
+
101
+ this.adapter = parsed.activeLearning?.enable
102
+ ? new TableAdapter({ client: this.client, tableName: parsed.activeLearning.tableName })
103
+ : new MemoryAdapter([])
104
+ }
105
+
106
+ /** @internal */
107
+ protected async callModel(
108
+ props: Partial<llm.generateContent.Input>
109
+ ): Promise<llm.generateContent.Output & { metadata: GenerationMetadata }> {
110
+ let retries = this.retry.maxRetries
111
+ while (retries-- >= 0) {
112
+ try {
113
+ return await this._callModel(props)
114
+ } catch (e) {
115
+ if (retries >= 0) {
116
+ await new Promise((resolve) => setTimeout(resolve, 1000))
117
+ } else {
118
+ throw new Error('Failed to call model after multiple retries')
119
+ }
120
+ }
121
+ }
122
+
123
+ throw new Error('Failed to call model after multiple retries')
124
+ }
125
+
126
+ /** @internal */
127
+ private async _callModel(
128
+ props: Partial<llm.generateContent.Input>
129
+ ): Promise<llm.generateContent.Output & { metadata: GenerationMetadata }> {
130
+ let retries = this.retry.maxRetries
131
+ do {
132
+ const start = Date.now()
133
+ const input: llm.generateContent.Input = {
134
+ messages: [],
135
+ temperature: 0.0,
136
+ topP: 1,
137
+ model: { id: this.model },
138
+ userId: this.userId,
139
+ ...props
140
+ }
141
+
142
+ const { output } = (await this.client.callAction({
143
+ type: `${this.integration}:generateContent`,
144
+ input
145
+ })) as unknown as { output: llm.generateContent.Output }
146
+
147
+ const latency = Date.now() - start
148
+
149
+ return {
150
+ ...output,
151
+ metadata: {
152
+ model: this.model,
153
+ latency,
154
+ cost: { input: output.usage.inputCost, output: output.usage.outputCost },
155
+ tokens: { input: output.usage.inputTokens, output: output.usage.outputTokens }
156
+ }
157
+ }
158
+ } while (--retries > 0)
159
+ }
160
+
161
+ protected async getTokenizer() {
162
+ Zai.tokenizer ??= await (async () => {
163
+ while (!getWasmTokenizer) {
164
+ // there's an issue with wasm, it doesn't load immediately
165
+ await new Promise((resolve) => setTimeout(resolve, 25))
166
+ }
167
+ return getWasmTokenizer()
168
+ })()
169
+ return Zai.tokenizer
170
+ }
171
+
172
+ protected get taskId() {
173
+ if (!this.activeLearning.enable) {
174
+ return undefined
175
+ }
176
+
177
+ return `${this.namespace}/${this.activeLearning.taskId}`.replace(/\/+/g, '/')
178
+ }
179
+
180
+ public with(options: Partial<ZaiConfig>): Zai {
181
+ return new Zai({
182
+ ...this.originalConfig,
183
+ ...options
184
+ })
185
+ }
186
+
187
+ public learn(taskId: string) {
188
+ return new Zai({
189
+ ...this.originalConfig,
190
+ activeLearning: { ...this.activeLearning, taskId, enable: true }
191
+ })
192
+ }
193
+ }
package/tsconfig.json CHANGED
@@ -3,6 +3,7 @@
3
3
  "target": "ESNext",
4
4
  "module": "ESNext",
5
5
  "moduleResolution": "bundler",
6
+ "outDir": "dist",
6
7
  "allowJs": true,
7
8
  "skipLibCheck": true,
8
9
  "esModuleInterop": true,
@@ -16,7 +17,7 @@
16
17
  "noUnusedParameters": true,
17
18
  "noUncheckedIndexedAccess": true,
18
19
  "lib": ["dom", "ESNext", "dom.iterable"],
19
- "declaration": false,
20
+ "declaration": true,
20
21
  "noEmit": false,
21
22
  "paths": {
22
23
  "@botpress/zai": ["./src/zai.ts"]
@@ -24,7 +25,6 @@
24
25
  },
25
26
  "exclude": ["node_modules", "dist"],
26
27
  "include": ["src/**/*", "vitest.d.ts"],
27
-
28
28
  "ts-node": {
29
29
  "esm": true,
30
30
  "require": ["dotenv/config", "./ensure-env.cjs"]