411 lines
13 KiB
TypeScript
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 },
|
|
})
|
|
}
|
|
}
|