2023-05-04 03:54:32 +00:00
|
|
|
import { ZodRawShape, ZodTypeAny, z } from 'zod'
|
2023-05-02 02:32:03 +00:00
|
|
|
|
|
|
|
import * as types from './types'
|
|
|
|
|
2023-05-04 03:54:32 +00:00
|
|
|
export abstract class BaseLLMCallBuilder<
|
|
|
|
TInput extends ZodRawShape | ZodTypeAny = ZodTypeAny,
|
2023-05-04 19:48:28 +00:00
|
|
|
TOutput extends ZodRawShape | ZodTypeAny = z.ZodType<string>,
|
2023-05-04 03:54:32 +00:00
|
|
|
TModelParams extends Record<string, any> = Record<string, any>
|
|
|
|
> {
|
2023-05-02 02:32:03 +00:00
|
|
|
_options: types.BaseLLMOptions<TInput, TOutput, TModelParams>
|
|
|
|
|
|
|
|
constructor(options: types.BaseLLMOptions<TInput, TOutput, TModelParams>) {
|
|
|
|
this._options = options
|
|
|
|
}
|
|
|
|
|
2023-05-04 03:54:32 +00:00
|
|
|
input<U extends ZodRawShape | ZodTypeAny = TInput>(
|
|
|
|
inputSchema: U
|
|
|
|
): BaseLLMCallBuilder<U, TOutput, TModelParams> {
|
|
|
|
;(
|
|
|
|
this as unknown as BaseLLMCallBuilder<U, TOutput, TModelParams>
|
|
|
|
)._options.input = inputSchema
|
|
|
|
return this as unknown as BaseLLMCallBuilder<U, TOutput, TModelParams>
|
2023-05-02 02:32:03 +00:00
|
|
|
}
|
|
|
|
|
2023-05-04 03:54:32 +00:00
|
|
|
output<U extends ZodRawShape | ZodTypeAny = TOutput>(
|
|
|
|
outputSchema: U
|
|
|
|
): BaseLLMCallBuilder<TInput, U, TModelParams> {
|
|
|
|
;(
|
|
|
|
this as unknown as BaseLLMCallBuilder<TInput, U, TModelParams>
|
|
|
|
)._options.output = outputSchema
|
|
|
|
return this as unknown as BaseLLMCallBuilder<TInput, U, TModelParams>
|
2023-05-02 02:32:03 +00:00
|
|
|
}
|
|
|
|
|
|
|
|
examples(examples: types.LLMExample[]) {
|
|
|
|
this._options.examples = examples
|
|
|
|
return this
|
|
|
|
}
|
|
|
|
|
2023-05-25 21:59:22 +00:00
|
|
|
modelParams(params: Partial<TModelParams>) {
|
|
|
|
// We assume that modelParams does not include nested objects; if it did, we would need to do a deep merge...
|
|
|
|
this._options.modelParams = Object.assign(
|
|
|
|
{},
|
|
|
|
this._options.modelParams,
|
|
|
|
params
|
|
|
|
)
|
|
|
|
return this
|
|
|
|
}
|
|
|
|
|
2023-05-02 02:32:03 +00:00
|
|
|
retry(retryConfig: types.LLMRetryConfig) {
|
|
|
|
this._options.retryConfig = retryConfig
|
|
|
|
return this
|
|
|
|
}
|
|
|
|
|
2023-05-04 03:54:32 +00:00
|
|
|
abstract call(
|
|
|
|
input?: types.ParsedData<TInput>
|
|
|
|
): Promise<types.ParsedData<TOutput>>
|
2023-05-02 06:44:08 +00:00
|
|
|
|
|
|
|
// TODO
|
2023-05-04 03:54:32 +00:00
|
|
|
// abstract stream({
|
2023-05-02 06:44:08 +00:00
|
|
|
// input: TInput,
|
|
|
|
// onProgress: types.ProgressFunction
|
2023-05-04 03:54:32 +00:00
|
|
|
// }): Promise<TOutput>
|
2023-05-02 02:32:03 +00:00
|
|
|
}
|
|
|
|
|
|
|
|
export abstract class ChatModelBuilder<
|
2023-05-04 03:54:32 +00:00
|
|
|
TInput extends ZodRawShape | ZodTypeAny = ZodTypeAny,
|
2023-05-04 19:48:28 +00:00
|
|
|
TOutput extends ZodRawShape | ZodTypeAny = z.ZodType<string>,
|
2023-05-04 03:54:32 +00:00
|
|
|
TModelParams extends Record<string, any> = Record<string, any>
|
2023-05-02 02:32:03 +00:00
|
|
|
> extends BaseLLMCallBuilder<TInput, TOutput, TModelParams> {
|
|
|
|
_messages: types.ChatMessage[]
|
|
|
|
|
|
|
|
constructor(options: types.ChatModelOptions<TInput, TOutput, TModelParams>) {
|
|
|
|
super(options)
|
|
|
|
|
|
|
|
this._messages = options.messages
|
|
|
|
}
|
|
|
|
}
|