151 lines
3.6 KiB
TypeScript
151 lines
3.6 KiB
TypeScript
import { z } from 'zod'
|
|
import { variableStringSchema } from '../../utils'
|
|
import { blockBaseSchema, credentialsBaseSchema } from '../baseSchemas'
|
|
import { IntegrationBlockType } from './enums'
|
|
|
|
export const openAITasks = ['Create chat completion', 'Create image'] as const
|
|
|
|
export const chatCompletionModels = [
|
|
'gpt-4',
|
|
'gpt-4-0314',
|
|
'gpt-4-32k',
|
|
'gpt-4-32k-0314',
|
|
'gpt-3.5-turbo',
|
|
'gpt-3.5-turbo-0301',
|
|
] as const
|
|
|
|
export const modelLimit = {
|
|
'gpt-3.5-turbo': 4096,
|
|
'gpt-3.5-turbo-0301': 4096,
|
|
'gpt-4': 8192,
|
|
'gpt-4-0314': 8192,
|
|
'gpt-4-32k': 32768,
|
|
'gpt-4-32k-0314': 32768,
|
|
} as const
|
|
|
|
export const chatCompletionMessageRoles = [
|
|
'system',
|
|
'user',
|
|
'assistant',
|
|
] as const
|
|
|
|
export const chatCompletionMessageCustomRoles = [
|
|
'Messages sequence ✨',
|
|
] as const
|
|
|
|
export const chatCompletionResponseValues = [
|
|
'Message content',
|
|
'Total tokens',
|
|
] as const
|
|
|
|
const openAIBaseOptionsSchema = z.object({
|
|
credentialsId: z.string().optional(),
|
|
})
|
|
|
|
const initialOptionsSchema = z
|
|
.object({
|
|
task: z.undefined(),
|
|
})
|
|
.merge(openAIBaseOptionsSchema)
|
|
|
|
const chatCompletionMessageSchema = z.object({
|
|
id: z.string(),
|
|
role: z.enum(chatCompletionMessageRoles).optional(),
|
|
content: z.string().optional(),
|
|
})
|
|
|
|
const chatCompletionCustomMessageSchema = z.object({
|
|
id: z.string(),
|
|
role: z.enum(chatCompletionMessageCustomRoles),
|
|
content: z
|
|
.object({
|
|
assistantMessagesVariableId: z.string().optional(),
|
|
userMessagesVariableId: z.string().optional(),
|
|
})
|
|
.optional(),
|
|
})
|
|
|
|
const chatCompletionOptionsSchema = z
|
|
.object({
|
|
task: z.literal(openAITasks[0]),
|
|
model: z.enum(chatCompletionModels),
|
|
messages: z.array(
|
|
z.union([chatCompletionMessageSchema, chatCompletionCustomMessageSchema])
|
|
),
|
|
advancedSettings: z
|
|
.object({
|
|
temperature: z.number().or(variableStringSchema).optional(),
|
|
})
|
|
.optional(),
|
|
responseMapping: z.array(
|
|
z.object({
|
|
id: z.string(),
|
|
valueToExtract: z.enum(chatCompletionResponseValues),
|
|
variableId: z.string().optional(),
|
|
})
|
|
),
|
|
})
|
|
.merge(openAIBaseOptionsSchema)
|
|
|
|
const createImageOptionsSchema = z
|
|
.object({
|
|
task: z.literal(openAITasks[1]),
|
|
prompt: z.string().optional(),
|
|
advancedOptions: z.object({
|
|
size: z.enum(['256x256', '512x512', '1024x1024']).optional(),
|
|
}),
|
|
responseMapping: z.array(
|
|
z.object({
|
|
id: z.string(),
|
|
valueToExtract: z.enum(['Image URL']),
|
|
variableId: z.string().optional(),
|
|
})
|
|
),
|
|
})
|
|
.merge(openAIBaseOptionsSchema)
|
|
|
|
export const openAIBlockSchema = blockBaseSchema.merge(
|
|
z.object({
|
|
type: z.enum([IntegrationBlockType.OPEN_AI]),
|
|
options: z.discriminatedUnion('task', [
|
|
initialOptionsSchema,
|
|
chatCompletionOptionsSchema,
|
|
createImageOptionsSchema,
|
|
]),
|
|
})
|
|
)
|
|
|
|
export const openAICredentialsSchema = z
|
|
.object({
|
|
type: z.literal('openai'),
|
|
data: z.object({
|
|
apiKey: z.string(),
|
|
}),
|
|
})
|
|
.merge(credentialsBaseSchema)
|
|
|
|
export const defaultChatCompletionOptions = (
|
|
createId: () => string
|
|
): ChatCompletionOpenAIOptions => ({
|
|
task: 'Create chat completion',
|
|
messages: [
|
|
{
|
|
id: createId(),
|
|
},
|
|
],
|
|
responseMapping: [
|
|
{
|
|
id: createId(),
|
|
valueToExtract: 'Message content',
|
|
},
|
|
],
|
|
model: 'gpt-3.5-turbo',
|
|
})
|
|
|
|
export type OpenAICredentials = z.infer<typeof openAICredentialsSchema>
|
|
export type OpenAIBlock = z.infer<typeof openAIBlockSchema>
|
|
export type ChatCompletionOpenAIOptions = z.infer<
|
|
typeof chatCompletionOptionsSchema
|
|
>
|
|
export type CreateImageOpenAIOptions = z.infer<typeof createImageOptionsSchema>
|