✨ Add "Generate variables" actions in AI blocks

Closes #1586
2024-06-18 12:13:00 +02:00
parent bec9cb68ca
commit 76fcf7ee93
25 changed files with 860 additions and 165 deletions
--- a/packages/forge/blocks/openai/shared/runGenerateVariables.ts
+++ b/packages/forge/blocks/openai/shared/runGenerateVariables.ts
@@ -0,0 +1,101 @@
+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)
+}