[BUGFIX] Custom endpoint and form state (huggingface#398)

frascuchon · web-flow · commit b5277b3560e1 · 2025-07-22T07:58:28.000+02:00
* store endpoint usage flag in DB

* read endpoint usage from DB info

* show selected model properly

* create process with endpoint-related info

* remove extra condition
diff --git a/src/features/add-column/form/execution-form.tsx b/src/features/add-column/form/execution-form.tsx
@@ -78,7 +78,7 @@ export const ExecutionForm = component$<SidebarProps>(
 
     const allModels = useContext<Model[]>(modelsContext);
 
-    let {
+    const {
       DEFAULT_MODEL,
       DEFAULT_MODEL_PROVIDER,
       modelEndpointEnabled,
@@ -103,7 +103,8 @@ export const ExecutionForm = component$<SidebarProps>(
     const selectedModelId = useSignal<string>('');
     const selectedProvider = useSignal<string>('');
 
-    const endpointURLSelected = useSignal(modelEndpointEnabled);
+    const enableCustomEndpoint = useSignal(modelEndpointEnabled);
+    const endpointURLSelected = useSignal(false);
 
     const onSelectedVariables = $((variables: { id: string }[]) => {
       columnsReferences.value = variables.map((v) => v.id);
@@ -115,11 +116,6 @@ export const ExecutionForm = component$<SidebarProps>(
       return column.type === 'image';
     });
 
-    if (isImageColumn.value) {
-      // Currently, we custom endpoint only for text models
-      modelEndpointEnabled = false;
-    }
-
     const modelProviders = useComputed$(() => {
       const model = models.value.find(
         (m: Model) =>
@@ -166,19 +162,28 @@ export const ExecutionForm = component$<SidebarProps>(
       }
     });
 
-    useTask$(() => {
+    useTask$(({ track }) => {
+      track(column);
+
       variables.value = columns.value
         .filter((c) => c.id !== column.id && !hasBlobContent(c))
         .map((c) => ({
           id: c.id,
           name: c.name,
         }));
 
+      if (isImageColumn.value) {
+        // Currently, we custom endpoint only for text models
+        enableCustomEndpoint.value = false;
+      }
+
       const { process } = column;
       if (!process) return;
 
       prompt.value = process.prompt;
       searchOnWeb.value = process.searchEnabled || false;
+      endpointURLSelected.value =
+        (enableCustomEndpoint.value && process.useEndpointURL) || false;
 
       if (process.modelName) {
         // If there's a previously selected model, use that
@@ -385,7 +390,7 @@ export const ExecutionForm = component$<SidebarProps>(
                     <p class="text-neutral-500 underline">
                       {selectedModelId.value}
                     </p>
-                    {modelEndpointEnabled && !endpointURLSelected.value && (
+                    {!endpointURLSelected.value && (
                       <Tooltip text="Reset default model">
                         <LuUndo2
                           class="w-4 h-4 rounded-full gap-2 text-neutral-500 cursor-pointer hover:bg-neutral-200"
diff --git a/src/features/table/components/header/cell-generation.tsx b/src/features/table/components/header/cell-generation.tsx
@@ -1,8 +1,7 @@
-import { component$, useContext } from '@builder.io/qwik';
+import { component$ } from '@builder.io/qwik';
 import { LuEgg } from '@qwikest/icons/lucide';
 import { Tooltip } from '~/components/ui/tooltip/tooltip';
 import { useGenerateColumn } from '~/features/execution';
-import { configContext } from '~/routes/home/layout';
 import { type Column, TEMPORAL_ID } from '~/state';
 
 export const CellGeneration = component$<{ column: Column }>(({ column }) => {
@@ -11,10 +10,6 @@ export const CellGeneration = component$<{ column: Column }>(({ column }) => {
   if (column.id === TEMPORAL_ID || column.kind !== 'dynamic') return null;
   if (!column.process) return null;
 
-  const { modelEndpointEnabled } = useContext(configContext);
-
-  column.process.useEndpointURL = modelEndpointEnabled;
-
   return (
     <Tooltip text="Regenerate">
       <div
diff --git a/src/features/table/table-body.tsx b/src/features/table/table-body.tsx
@@ -219,7 +219,6 @@ export const TableBody = component$(() => {
       ...column,
       process: {
         ...column.process!,
-        useEndpointURL: modelEndpointEnabled,
         offset,
         limit,
       },
diff --git a/src/services/db/models/process.ts b/src/services/db/models/process.ts
@@ -19,8 +19,11 @@ export class ProcessModel extends Model<
   declare id: CreationOptional<string>;
   declare prompt: string;
   declare modelName: string;
-  declare searchEnabled: boolean;
+
   declare modelProvider: string;
+  declare useCustomEndpoint: boolean;
+
+  declare searchEnabled: boolean;
   declare columnId: ForeignKey<ColumnModel['id']>;
 
   declare referredColumns: NonAttribute<ColumnModel[]>; // This is a virtual attribute
@@ -40,18 +43,22 @@ ProcessModel.init(
       defaultValue: DataTypes.UUIDV4,
       primaryKey: true,
     },
-    modelName: {
+    prompt: {
       type: DataTypes.STRING,
       allowNull: false,
     },
-    modelProvider: {
+    modelName: {
       type: DataTypes.STRING,
       allowNull: false,
     },
-    prompt: {
+    modelProvider: {
       type: DataTypes.STRING,
       allowNull: false,
     },
+    useCustomEndpoint: {
+      type: DataTypes.BOOLEAN,
+      defaultValue: false,
+    },
     searchEnabled: {
       type: DataTypes.BOOLEAN,
       defaultValue: false,
diff --git a/src/services/repository/columns.ts b/src/services/repository/columns.ts
@@ -26,9 +26,10 @@ export const modelToColumn = (model: ColumnModel): Column => {
       columnsReferences: (model.process?.referredColumns ?? []).map(
         (columnRef) => columnRef.id,
       ),
+      prompt: model.process?.prompt ?? '',
       modelName: model.process?.modelName ?? '',
       modelProvider: model.process?.modelProvider ?? '',
-      prompt: model.process?.prompt ?? '',
+      useEndpointURL: model.process?.useCustomEndpoint ?? false,
       searchEnabled: model.process?.searchEnabled,
       updatedAt: model.process?.updatedAt,
     },
diff --git a/src/services/repository/processes.ts b/src/services/repository/processes.ts
@@ -2,9 +2,10 @@ import { ProcessColumnModel, ProcessModel } from '~/services/db/models';
 import type { Process } from '~/state';
 
 export interface CreateProcess {
+  prompt: string;
   modelName: string;
   modelProvider: string;
-  prompt: string;
+  useEndpointURL?: boolean;
   searchEnabled: boolean;
   columnsReferences?: string[];
 }
@@ -19,9 +20,10 @@ export const createProcess = async ({
   };
 }): Promise<Process> => {
   const model = await ProcessModel.create({
+    prompt: process.prompt,
     modelName: process.modelName,
     modelProvider: process.modelProvider,
-    prompt: process.prompt,
+    useCustomEndpoint: process.useEndpointURL ?? false,
     searchEnabled: process.searchEnabled,
     columnId: column.id,
   });
@@ -37,9 +39,10 @@ export const createProcess = async ({
 
   return {
     id: model.id,
+    prompt: model.prompt,
     modelName: model.modelName,
     modelProvider: model.modelProvider,
-    prompt: model.prompt,
+    useEndpointURL: model.useCustomEndpoint,
     searchEnabled: model.searchEnabled,
     columnsReferences: process?.columnsReferences || [],
     updatedAt: model.updatedAt,
@@ -55,9 +58,10 @@ export const updateProcess = async (process: Process): Promise<Process> => {
 
   model.changed('updatedAt', true);
   model.set({
+    prompt: process.prompt,
     modelName: process.modelName,
     modelProvider: process.modelProvider,
-    prompt: process.prompt,
+    useCustomEndpoint: process.useEndpointURL ?? false,
     searchEnabled: process.searchEnabled,
   });
 
@@ -74,9 +78,10 @@ export const updateProcess = async (process: Process): Promise<Process> => {
 
   return {
     id: model.id,
+    prompt: model.prompt,
     modelName: model.modelName,
     modelProvider: model.modelProvider,
-    prompt: model.prompt,
+    useEndpointURL: model.useCustomEndpoint,
     searchEnabled: model.searchEnabled,
     columnsReferences: process.columnsReferences,
     updatedAt: model.updatedAt,
diff --git a/src/usecases/run-autodataset.ts b/src/usecases/run-autodataset.ts
@@ -385,9 +385,10 @@ async function createDatasetWithColumns(
 
     const process = await createProcess({
       process: {
+        prompt: column.prompt,
         modelName: processModelName,
         modelProvider: processModelProvider,
-        prompt: column.prompt,
+        useEndpointURL: textGeneration.endpointUrl !== undefined && !isImage,
         searchEnabled,
         columnsReferences: columnReferences.map((ref) => {
           const refIndex = columnNames.indexOf(ref);
