102 lines
2.9 KiB
TypeScript
102 lines
2.9 KiB
TypeScript
import { LogsStore, VariableStore } from '@typebot.io/forge/types'
|
|
import {
|
|
GenerateVariablesOptions,
|
|
toolParametersSchema,
|
|
} from './parseGenerateVariablesOptions'
|
|
import { generateObject, LanguageModel } from 'ai'
|
|
import { Variable } from '@typebot.io/variables/types'
|
|
import { z } from '@typebot.io/forge/zod'
|
|
import { isNotEmpty } from '@typebot.io/lib/utils'
|
|
|
|
type Props = {
|
|
model: LanguageModel
|
|
variables: VariableStore
|
|
logs: LogsStore
|
|
} & Pick<GenerateVariablesOptions, 'variablesToExtract' | 'prompt'>
|
|
|
|
export const runGenerateVariables = async ({
|
|
variablesToExtract,
|
|
model,
|
|
prompt,
|
|
variables: variablesStore,
|
|
logs,
|
|
}: Props) => {
|
|
if (!prompt) return logs.add('No prompt provided')
|
|
const variables = variablesStore.list()
|
|
|
|
const schema = convertVariablesToExtractToSchema({
|
|
variablesToExtract,
|
|
variables,
|
|
})
|
|
if (!schema) {
|
|
logs.add('Could not parse variables to extract')
|
|
return
|
|
}
|
|
|
|
const hasOptionalVariables = variablesToExtract?.some(
|
|
(variableToExtract) => variableToExtract.isRequired === false
|
|
)
|
|
|
|
const { object } = await generateObject({
|
|
model,
|
|
schema,
|
|
prompt:
|
|
`${prompt}\n\nYou should generate a JSON object` +
|
|
(hasOptionalVariables
|
|
? ' and provide empty values if the information is not there or if you are unsure.'
|
|
: '.'),
|
|
})
|
|
|
|
Object.entries(object).forEach(([key, value]) => {
|
|
if (value === null) return
|
|
const existingVariable = variables.find((v) => v.name === key)
|
|
if (!existingVariable) return
|
|
variablesStore.set(existingVariable.id, value)
|
|
})
|
|
}
|
|
|
|
const convertVariablesToExtractToSchema = ({
|
|
variablesToExtract,
|
|
variables,
|
|
}: {
|
|
variablesToExtract: z.infer<typeof toolParametersSchema> | undefined
|
|
variables: Variable[]
|
|
}): z.ZodTypeAny | undefined => {
|
|
if (!variablesToExtract || variablesToExtract?.length === 0) return
|
|
|
|
const shape: z.ZodRawShape = {}
|
|
variablesToExtract.forEach((variableToExtract) => {
|
|
if (!variableToExtract) return
|
|
const matchingVariable = variables.find(
|
|
(v) => v.id === variableToExtract.variableId
|
|
)
|
|
if (!matchingVariable) return
|
|
switch (variableToExtract.type) {
|
|
case 'string':
|
|
shape[matchingVariable.name] = z.string()
|
|
break
|
|
case 'number':
|
|
shape[matchingVariable.name] = z.number()
|
|
break
|
|
case 'boolean':
|
|
shape[matchingVariable.name] = z.boolean()
|
|
break
|
|
case 'enum': {
|
|
if (!variableToExtract.values || variableToExtract.values.length === 0)
|
|
return
|
|
shape[matchingVariable.name] = z.enum(variableToExtract.values as any)
|
|
break
|
|
}
|
|
}
|
|
if (variableToExtract.isRequired === false)
|
|
shape[matchingVariable.name] = shape[matchingVariable.name].optional()
|
|
|
|
if (isNotEmpty(variableToExtract.description))
|
|
shape[matchingVariable.name] = shape[matchingVariable.name].describe(
|
|
variableToExtract.description
|
|
)
|
|
})
|
|
|
|
return z.object(shape)
|
|
}
|