feat: dashboard: add model field to auto-embeddings form (#2661)

fixes #2660
nhost · Apr 15, 2024 · dd5d262 · dd5d262
1 parent 09962be
commit dd5d262
Show file tree

Hide file tree

Showing 6 changed files with 518 additions and 623 deletions.
diff --git a/.changeset/quick-boats-rhyme.md b/.changeset/quick-boats-rhyme.md
@@ -0,0 +1,5 @@
+---
+'@nhost/dashboard': minor
+---
+
+feat: add model field to the auto-embeddings form
diff --git a/dashboard/src/features/ai/AutoEmbeddingsForm/AutoEmbeddingsForm.tsx b/dashboard/src/features/ai/AutoEmbeddingsForm/AutoEmbeddingsForm.tsx
@@ -1,11 +1,13 @@
 import { useDialog } from '@/components/common/DialogProvider';
+import { ControlledSelect } from '@/components/form/ControlledSelect';
 import { Form } from '@/components/form/Form';
 import { Box } from '@/components/ui/v2/Box';
 import { Button } from '@/components/ui/v2/Button';
 import { ArrowsClockwise } from '@/components/ui/v2/icons/ArrowsClockwise';
 import { InfoIcon } from '@/components/ui/v2/icons/InfoIcon';
 import { PlusIcon } from '@/components/ui/v2/icons/PlusIcon';
 import { Input } from '@/components/ui/v2/Input';
+import { Option } from '@/components/ui/v2/Option';
 import { Text } from '@/components/ui/v2/Text';
 import { Tooltip } from '@/components/ui/v2/Tooltip';
 import { useAdminApolloClient } from '@/features/projects/common/hooks/useAdminApolloClient';
@@ -20,11 +22,18 @@ import { useEffect } from 'react';
 import { FormProvider, useForm } from 'react-hook-form';
 import * as Yup from 'yup';
 
+const AUTO_EMBEDDINGS_MODELS = [
+  'text-embedding-ada-002',
+  'text-embedding-3-small',
+  'text-embedding-3-large',
+];
+
 export const validationSchema = Yup.object({
-  name: Yup.string().required('The name is required.'),
-  schemaName: Yup.string().required('The schema is required'),
-  tableName: Yup.string().required('The table is required'),
-  columnName: Yup.string().required('The column is required'),
+  name: Yup.string().required('The name field is required.'),
+  model: Yup.string().oneOf(AUTO_EMBEDDINGS_MODELS),
+  schemaName: Yup.string().required('The schema field is required'),
+  tableName: Yup.string().required('The table field is required'),
+  columnName: Yup.string().required('The column field is required'),
   query: Yup.string(),
   mutation: Yup.string(),
 });
@@ -40,7 +49,7 @@ export interface AutoEmbeddingsFormProps extends DialogFormProps {
   /**
    * if there is initialData then it's an update operation
    */
-  initialData?: AutoEmbeddingsFormValues;
+  initialData?: AutoEmbeddingsFormValues & { model: string };
 
   /**
    * Function to be called when the operation is cancelled.
@@ -74,7 +83,10 @@ export default function AutoEmbeddingsForm({
     });
 
   const form = useForm<AutoEmbeddingsFormValues>({
-    defaultValues: initialData,
+    defaultValues: {
+      ...initialData,
+      model: initialData?.model ?? 'text-embedding-ada-002',
+    },
     reValidateMode: 'onSubmit',
     resolver: yupResolver(validationSchema),
   });
@@ -106,7 +118,9 @@ export default function AutoEmbeddingsForm({
     }
 
     await insertGraphiteAutoEmbeddingsConfiguration({
-      variables: values,
+      variables: {
+        ...values,
+      },
     });
   };
 
@@ -129,9 +143,9 @@ export default function AutoEmbeddingsForm({
     <FormProvider {...form}>
       <Form
         onSubmit={handleSubmit}
-        className="flex h-full flex-col gap-4 overflow-hidden"
+        className="flex flex-col h-full gap-4 overflow-hidden"
       >
-        <div className="flex flex-1 flex-col space-y-6 overflow-auto px-6">
+        <div className="flex flex-col flex-1 px-6 space-y-6 overflow-auto">
           <Input
             {...register('name')}
             id="name"
@@ -141,7 +155,7 @@ export default function AutoEmbeddingsForm({
                 <Tooltip title="Name of the Auto-Embeddings">
                   <InfoIcon
                     aria-label="Info"
-                    className="h-4 w-4"
+                    className="w-4 h-4"
                     color="primary"
                   />
                 </Tooltip>
@@ -155,6 +169,36 @@ export default function AutoEmbeddingsForm({
             autoComplete="off"
             autoFocus
           />
+
+          <ControlledSelect
+            slotProps={{
+              popper: { disablePortal: false, className: 'z-[10000]' },
+            }}
+            id="model"
+            name="model"
+            label={
+              <Box className="flex flex-row items-center space-x-2">
+                <Text>Model</Text>
+                <Tooltip title="Auto-Embeddings Model">
+                  <InfoIcon
+                    aria-label="Info"
+                    className="w-4 h-4"
+                    color="primary"
+                  />
+                </Tooltip>
+              </Box>
+            }
+            fullWidth
+            error={!!errors?.model?.message}
+            helperText={errors?.model?.message}
+          >
+            {AUTO_EMBEDDINGS_MODELS.map((model) => (
+              <Option key={model} value={model}>
+                {model}
+              </Option>
+            ))}
+          </ControlledSelect>
+
           <Input
             {...register('schemaName')}
             id="schemaName"
@@ -164,7 +208,7 @@ export default function AutoEmbeddingsForm({
                 <Tooltip title={<span>Schema where the table belongs to</span>}>
                   <InfoIcon
                     aria-label="Info"
-                    className="h-4 w-4"
+                    className="w-4 h-4"
                     color="primary"
                   />
                 </Tooltip>
@@ -186,7 +230,7 @@ export default function AutoEmbeddingsForm({
                 <Tooltip title="Table Name">
                   <InfoIcon
                     aria-label="Info"
-                    className="h-4 w-4"
+                    className="w-4 h-4"
                     color="primary"
                   />
                 </Tooltip>
@@ -208,7 +252,7 @@ export default function AutoEmbeddingsForm({
                 <Tooltip title="Column name">
                   <InfoIcon
                     aria-label="Info"
-                    className="h-4 w-4"
+                    className="w-4 h-4"
                     color="primary"
                   />
                 </Tooltip>
@@ -230,7 +274,7 @@ export default function AutoEmbeddingsForm({
                 <Tooltip title="Query">
                   <InfoIcon
                     aria-label="Info"
-                    className="h-4 w-4"
+                    className="w-4 h-4"
                     color="primary"
                   />
                 </Tooltip>
@@ -254,7 +298,7 @@ export default function AutoEmbeddingsForm({
                 <Tooltip title="Mutation">
                   <InfoIcon
                     aria-label="Info"
-                    className="h-4 w-4"
+                    className="w-4 h-4"
                     color="primary"
                   />
                 </Tooltip>
@@ -271,7 +315,7 @@ export default function AutoEmbeddingsForm({
           />
         </div>
 
-        <Box className="flex w-full flex-row justify-between rounded border-t px-6 py-4">
+        <Box className="flex flex-row justify-between w-full px-6 py-4 border-t rounded">
           <Button variant="outlined" color="secondary" onClick={onCancel}>
             Cancel
           </Button>

diff --git a/dashboard/src/gql/graphite/auto-embeddings/getGraphiteAutoEmbeddingsConfigurations.gql b/dashboard/src/gql/graphite/auto-embeddings/getGraphiteAutoEmbeddingsConfigurations.gql
@@ -2,6 +2,7 @@ query getGraphiteAutoEmbeddingsConfigurations($limit: Int!, $offset: Int!) {
   graphiteAutoEmbeddingsConfigurations(limit: $limit, offset: $offset) {
     id
     name
+    model
     schemaName
     tableName
     columnName

diff --git a/dashboard/src/gql/graphite/auto-embeddings/insertGraphiteAutoEmbeddingsConfiguration.gql b/dashboard/src/gql/graphite/auto-embeddings/insertGraphiteAutoEmbeddingsConfiguration.gql
@@ -1,5 +1,6 @@
 mutation insertGraphiteAutoEmbeddingsConfiguration(
   $name: String
+  $model: embedding_model_enum
   $schemaName: String
   $tableName: String
   $columnName: String
@@ -9,6 +10,7 @@ mutation insertGraphiteAutoEmbeddingsConfiguration(
   insertGraphiteAutoEmbeddingsConfiguration(
     object: {
       name: $name
+      model: $model
       schemaName: $schemaName
       tableName: $tableName
       columnName: $columnName

diff --git a/dashboard/src/gql/graphite/auto-embeddings/updateGraphiteAutoEmbeddingsConfiguration.gql b/dashboard/src/gql/graphite/auto-embeddings/updateGraphiteAutoEmbeddingsConfiguration.gql
@@ -1,6 +1,7 @@
 mutation updateGraphiteAutoEmbeddingsConfiguration(
   $id: uuid!
   $name: String
+  $model: embedding_model_enum
   $schemaName: String
   $tableName: String
   $columnName: String
@@ -11,6 +12,7 @@ mutation updateGraphiteAutoEmbeddingsConfiguration(
     pk_columns: { id: $id }
     _set: {
       name: $name
+      model: $model
       schemaName: $schemaName
       tableName: $tableName
       columnName: $columnName
@@ -20,6 +22,7 @@ mutation updateGraphiteAutoEmbeddingsConfiguration(
   ) {
     id
     name
+    model
     schemaName
     tableName
     columnName