Files
botpress--botpress/packages/zai/src/zai.ts
T
2026-07-13 13:34:48 +08:00

411 lines
13 KiB
TypeScript

import { Client } from '@botpress/client'
import { BotpressClientLike, Cognitive, Model, Models } from '@botpress/cognitive'
import { z } from '@bpinternal/zui'
import { Adapter } from './adapters/adapter'
import { TableAdapter } from './adapters/botpress-table'
import { MemoryAdapter } from './adapters/memory'
import { type TextTokenizer, getTokenizer } from './tokenizer'
/**
* A memoizer that caches the result of async operations by a unique key.
*
* When used with the Botpress ADK workflow `step` function, this enables
* Zai operations to resume where they left off if a workflow is interrupted.
*
*/
export type Memoizer = {
run: <T>(id: string, fn: () => Promise<T>) => Promise<T>
}
/**
* Active learning configuration for improving AI operations over time.
*
* When enabled, Zai stores successful operation results in a table and uses them as examples
* for future operations, improving accuracy and consistency.
*
* @example
* ```typescript
* const activeLearning = {
* enable: true,
* tableName: 'MyAppLearningTable',
* taskId: 'sentiment-analysis'
* }
* ```
*/
type ActiveLearning = {
/** Whether to enable active learning for this Zai instance */
enable: boolean
/** Name of the Botpress table to store learning examples (must end with 'Table') */
tableName: string
/** Unique identifier for this learning task */
taskId: string
}
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'),
})
/**
* Configuration options for creating a Zai instance.
*
* @example
* ```typescript
* import { Client } from '@botpress/client'
* import { Zai } from '@botpress/zai'
*
* const client = new Client({ token: 'your-token' })
* const config: ZaiConfig = {
* client,
* modelId: 'best', // Use the best available model
* userId: 'user-123',
* namespace: 'my-app',
* activeLearning: {
* enable: true,
* tableName: 'MyLearningTable',
* taskId: 'extraction'
* }
* }
* const zai = new Zai(config)
* ```
*/
export type ZaiConfig = {
/** Botpress client or Cognitive client instance */
client: BotpressClientLike | Cognitive
/** Optional user ID for tracking and attribution */
userId?: string
/**
* Model to use: 'best' (default), 'fast', or specific model like 'openai:gpt-4'.
* An array can be provided to specify ordered fallback models if the primary
* model is unavailable. Server-side fallback is honored on the cognitive-v2
* path; the legacy integration path uses the first entry only.
*/
modelId?: Models | Models[]
/** Active learning configuration to improve operations over time */
activeLearning?: ActiveLearning
/** Namespace for organizing tasks (default: 'zai') */
namespace?: string
/**
* Memoizer (or factory returning one) for caching cognitive call results.
*
* When provided, all LLM calls are wrapped in the memoizer, allowing results
* to be cached and replayed. This is useful for resuming workflow runs where
* Zai operations have already completed their cognitive calls.
*
* If a factory function is provided, it is called once per Zai operation invocation.
*/
memoize?: Memoizer | (() => Memoizer)
}
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<Models | Models[]>(
(value) => {
const isValidSingle = (v: unknown): v is string =>
typeof v === 'string' && (v === 'best' || v === 'fast' || v === 'auto' || v.includes(':'))
if (Array.isArray(value)) {
return value.length > 0 && value.every(isValidSingle)
}
return isValidSingle(value)
},
{
message: 'At least one model ID is invalid. Expected a model string or an array of model strings.',
}
)
.describe('The ID of the model you want to use, or an ordered list of fallback models')
.default('best' satisfies Models),
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'),
})
/**
* Zai - A type-safe LLM utility library for production-ready AI operations.
*
* Zai provides high-level abstractions for common AI tasks with built-in features like:
* - Active learning (learns from successful operations)
* - Automatic chunking for large inputs
* - Retry logic with error recovery
* - Usage tracking (tokens, cost, latency)
* - Type-safe schema validation with Zod
*
* @example Basic usage
* ```typescript
* import { Client } from '@botpress/client'
* import { Zai } from '@botpress/zai'
* import { z } from '@bpinternal/zui'
*
* const client = new Client({ token: process.env.BOTPRESS_TOKEN })
* const zai = new Zai({ client })
*
* // Extract structured data
* const schema = z.object({
* name: z.string(),
* age: z.number()
* })
* const person = await zai.extract('John is 30 years old', schema)
* // Output: { name: 'John', age: 30 }
*
* // Check conditions
* const isPositive = await zai.check('I love this product!', 'Is the sentiment positive?')
* // Output: true
*
* // Summarize text
* const summary = await zai.summarize(longDocument, { length: 100 })
* ```
*
* @example With active learning
* ```typescript
* const zai = new Zai({
* client,
* activeLearning: {
* enable: true,
* tableName: 'SentimentTable',
* taskId: 'product-reviews'
* }
* })
*
* // Enable learning for specific task
* const result = await zai.learn('sentiment').check(review, 'Is this positive?')
* // Future calls will use approved examples for better accuracy
* ```
*
* @example Chaining configuration
* ```typescript
* // Use fast model for quick operations
* const fastZai = zai.with({ modelId: 'fast' })
* await fastZai.check(text, 'Is this spam?')
*
* // Use specific model
* const gpt4Zai = zai.with({ modelId: 'openai:gpt-4' })
* await gpt4Zai.extract(document, complexSchema)
* ```
*/
export class Zai {
protected static tokenizer: TextTokenizer = null!
protected client: Cognitive
private _originalConfig: ZaiConfig
private _userId: string | undefined
protected Model: Models | Models[]
protected ModelDetails: Model
protected namespace: string
protected adapter: Adapter
protected activeLearning: ActiveLearning
protected _memoize?: Memoizer | (() => Memoizer)
/**
* Creates a new Zai instance with the specified configuration.
*
* @param config - Configuration object containing client, model, and learning settings
*
* @example
* ```typescript
* import { Client } from '@botpress/client'
* import { Zai } from '@botpress/zai'
*
* const client = new Client({ token: 'your-token' })
* const zai = new Zai({
* client,
* modelId: 'best',
* namespace: 'my-app',
* userId: 'user-123'
* })
* ```
*
* @throws {Error} If the configuration is invalid (e.g., invalid modelId format)
*/
public constructor(config: ZaiConfig) {
this._originalConfig = config
const parsed = _ZaiConfig.parse(config) as ZaiConfig
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 Models | Models[]
this.activeLearning = parsed.activeLearning as ActiveLearning
this.adapter = parsed.activeLearning?.enable
? new TableAdapter({
client: this.client.client as unknown as Client,
tableName: parsed.activeLearning.tableName,
})
: new MemoryAdapter([])
this._memoize = config.memoize
}
/** @internal */
protected async callModel(
props: Parameters<Cognitive['generateContent']>[0]
): ReturnType<Cognitive['generateContent']> {
return this.client.generateContent({
reasoningEffort: 'none',
...props,
model: this.Model as Required<Parameters<Cognitive['generateContent']>[0]>['model'],
userId: this._userId,
})
}
/** @internal */
protected _resolveMemoizer(): Memoizer | undefined {
if (!this._memoize) {
return undefined
}
return typeof this._memoize === 'function' ? this._memoize() : this._memoize
}
protected async getTokenizer() {
Zai.tokenizer ??= await getTokenizer()
return Zai.tokenizer
}
protected async fetchModelDetails(): Promise<void> {
if (!this.ModelDetails) {
// getModelDetails resolves a single model. When a fallback array is
// configured, we describe the primary model — fallbacks are only relevant
// at the request layer.
const primaryModel = Array.isArray(this.Model) ? this.Model[0] : this.Model
this.ModelDetails = await this.client.getModelDetails(primaryModel)
}
}
protected get taskId() {
if (!this.activeLearning.enable) {
return undefined
}
return `${this.namespace}/${this.activeLearning.taskId}`.replace(/\/+/g, '/')
}
/**
* Creates a new Zai instance with merged configuration options.
*
* This method allows you to create variations of your Zai instance with different
* settings without modifying the original. Useful for switching models, namespaces,
* or other configuration on a per-operation basis.
*
* @param options - Partial configuration to override the current settings
* @returns A new Zai instance with the merged configuration
*
* @example Switch to a faster model
* ```typescript
* const zai = new Zai({ client })
*
* // Use fast model for simple operations
* const fastZai = zai.with({ modelId: 'fast' })
* await fastZai.check(text, 'Is this spam?')
*
* // Use best model for complex operations
* const bestZai = zai.with({ modelId: 'best' })
* await bestZai.extract(document, complexSchema)
* ```
*
* @example Change namespace
* ```typescript
* const customerZai = zai.with({ namespace: 'customer-support' })
* const salesZai = zai.with({ namespace: 'sales' })
* ```
*
* @example Use specific model
* ```typescript
* const gpt4 = zai.with({ modelId: 'openai:gpt-4' })
* const claude = zai.with({ modelId: 'anthropic:claude-3-5-sonnet-20241022' })
* ```
*/
public with(options: Partial<ZaiConfig>): Zai {
return new Zai({
...this._originalConfig,
...options,
})
}
/**
* Creates a new Zai instance with active learning enabled for a specific task.
*
* Active learning stores successful operation results and uses them as examples for
* future operations, improving accuracy and consistency over time. Each task ID
* maintains its own set of learned examples.
*
* @param taskId - Unique identifier for the learning task (alphanumeric, hyphens, underscores, slashes)
* @returns A new Zai instance with active learning enabled for the specified task
*
* @example Sentiment analysis with learning
* ```typescript
* const zai = new Zai({
* client,
* activeLearning: {
* enable: false,
* tableName: 'AppLearningTable',
* taskId: 'default'
* }
* })
*
* // Enable learning for sentiment analysis
* const sentimentZai = zai.learn('sentiment-analysis')
* const result = await sentimentZai.check(review, 'Is this review positive?')
*
* // Each successful call is stored and used to improve future calls
* ```
*
* @example Different tasks for different purposes
* ```typescript
* // Extract user info with learning
* const userExtractor = zai.learn('user-extraction')
* await userExtractor.extract(text, userSchema)
*
* // Extract product info with separate learning
* const productExtractor = zai.learn('product-extraction')
* await productExtractor.extract(text, productSchema)
*
* // Each task learns independently
* ```
*
* @example Combining with other configuration
* ```typescript
* // Use fast model + learning
* const fastLearner = zai.with({ modelId: 'fast' }).learn('quick-checks')
* await fastLearner.check(email, 'Is this spam?')
* ```
*
* @see {@link ZaiConfig.activeLearning} for configuration options
*/
public learn(taskId: string) {
return new Zai({
...this._originalConfig,
activeLearning: { ...this.activeLearning, taskId, enable: true },
})
}
}