[OPIK-3263] [FE] Add model selector for OpikAI features (Prompt Generator, Prompt Improver) (#4572)

itamargolan · web-flow · commit 3d4a88ea2aec · 2025-12-29T17:54:06.000+02:00
* add a new hook for model selection and reuse it in dataset expansion and prompt improvement

* label style

* review fix

* review fixes - Sasha

* review fix
diff --git a/apps/opik-frontend/src/components/pages/DatasetItemsPage/DatasetExpansionDialog.tsx b/apps/opik-frontend/src/components/pages/DatasetItemsPage/DatasetExpansionDialog.tsx
@@ -31,9 +31,7 @@ import TooltipWrapper from "@/components/shared/TooltipWrapper/TooltipWrapper";
 import useDatasetExpansionMutation from "@/api/datasets/useDatasetExpansionMutation";
 import useDatasetItemsList from "@/api/datasets/useDatasetItemsList";
 import useAppStore from "@/store/AppStore";
-import useLastPickedModel from "@/hooks/useLastPickedModel";
-import useLLMProviderModelsData from "@/hooks/useLLMProviderModelsData";
-import useProviderKeys from "@/api/provider-keys/useProviderKeys";
+import useModelSelection from "@/hooks/useModelSelection";
 import useProgressSimulation from "@/hooks/useProgressSimulation";
 
 const DATASET_EXPANSION_PROGRESS_MESSAGES = [
@@ -43,7 +41,6 @@ const DATASET_EXPANSION_PROGRESS_MESSAGES = [
   "Finalizing generated data...",
 ];
 import { DatasetExpansionRequest, DatasetItem } from "@/types/datasets";
-import { COMPOSED_PROVIDER_TYPE, PROVIDER_MODEL_TYPE } from "@/types/providers";
 
 const DATASET_EXPANSION_LAST_PICKED_MODEL = "opik-dataset-expansion-model";
 const SAMPLE_COUNT_MIN = 1;
@@ -65,57 +62,11 @@ const DatasetExpansionDialog: React.FunctionComponent<
 > = ({ datasetId: initialDatasetId, open, setOpen, onSamplesGenerated }) => {
   const workspaceName = useAppStore((state) => state.activeWorkspaceName);
 
-  const [lastPickedModel, setLastPickedModel] = useLastPickedModel({
-    key: DATASET_EXPANSION_LAST_PICKED_MODEL,
+  // Model selection with persistence using the reusable hook
+  const { model, modelSelectProps } = useModelSelection({
+    persistenceKey: DATASET_EXPANSION_LAST_PICKED_MODEL,
   });
 
-  const { data: providerKeysData } = useProviderKeys({
-    workspaceName,
-  });
-
-  const providerKeys = useMemo(() => {
-    return providerKeysData?.content?.map((c) => c.ui_composed_provider) || [];
-  }, [providerKeysData]);
-
-  const { calculateModelProvider, calculateDefaultModel } =
-    useLLMProviderModelsData();
-
-  const { model, provider } = useMemo(() => {
-    const calculatedModel = calculateDefaultModel(
-      lastPickedModel,
-      providerKeys,
-    ) as PROVIDER_MODEL_TYPE;
-    const calculatedProvider = calculateModelProvider(calculatedModel);
-    return {
-      model: calculatedModel,
-      provider: calculatedProvider,
-    };
-  }, [
-    calculateDefaultModel,
-    calculateModelProvider,
-    lastPickedModel,
-    providerKeys,
-  ]);
-
-  const handleAddProvider = useCallback(
-    (provider: COMPOSED_PROVIDER_TYPE) => {
-      if (!model) {
-        setLastPickedModel(calculateDefaultModel(model, [provider], provider));
-      }
-    },
-    [calculateDefaultModel, model, setLastPickedModel],
-  );
-
-  const handleDeleteProvider = useCallback(
-    (provider: COMPOSED_PROVIDER_TYPE) => {
-      const currentProvider = calculateModelProvider(model);
-      if (currentProvider === provider) {
-        setLastPickedModel("");
-      }
-    },
-    [calculateModelProvider, model, setLastPickedModel],
-  );
-
   const [sampleCount, setSampleCount] = useState<number>(5);
   const [variationInstructions, setVariationInstructions] =
     useState<string>("");
@@ -325,11 +276,6 @@ const DatasetExpansionDialog: React.FunctionComponent<
     complete,
   ]);
 
-  const handleModelChange = useCallback(
-    (model: PROVIDER_MODEL_TYPE) => setLastPickedModel(model),
-    [setLastPickedModel],
-  );
-
   return (
     <Dialog open={open} onOpenChange={setOpen}>
       <DialogContent className="max-w-2xl">
@@ -539,12 +485,9 @@ const DatasetExpansionDialog: React.FunctionComponent<
           <div className="mt-6 space-y-2">
             <Label htmlFor="model">Model</Label>
             <PromptModelSelect
-              value={model}
               workspaceName={workspaceName}
-              onChange={handleModelChange}
-              provider={provider}
-              onAddProvider={handleAddProvider}
-              onDeleteProvider={handleDeleteProvider}
+              {...modelSelectProps}
+              disabled={isPending}
             />
           </div>
 
diff --git a/apps/opik-frontend/src/components/shared/PromptImprovementDialog/PromptImprovementDialog.tsx b/apps/opik-frontend/src/components/shared/PromptImprovementDialog/PromptImprovementDialog.tsx
@@ -22,8 +22,10 @@ import { Textarea } from "@/components/ui/textarea";
 import { Description } from "@/components/ui/description";
 import { Separator } from "@/components/ui/separator";
 import ExplainerDescription from "@/components/shared/ExplainerDescription/ExplainerDescription";
+import PromptModelSelect from "@/components/pages-shared/llm/PromptModelSelect/PromptModelSelect";
 import usePromptImprovement from "@/hooks/usePromptImprovement";
 import useProgressSimulation from "@/hooks/useProgressSimulation";
+import useModelSelection from "@/hooks/useModelSelection";
 import {
   COMPOSED_PROVIDER_TYPE,
   LLMPromptConfigsType,
@@ -49,6 +51,8 @@ const PROMPT_IMPROVEMENT_PROGRESS_MESSAGES = [
   "Polishing the output...",
 ];
 
+const PROMPT_IMPROVEMENT_LAST_PICKED_MODEL = "opik-prompt-improvement-model";
+
 interface PromptImprovementDialogProps {
   open: boolean;
   setOpen: (open: boolean) => void;
@@ -66,9 +70,9 @@ const PromptImprovementDialog: React.FC<PromptImprovementDialogProps> = ({
   setOpen,
   id,
   originalPrompt = "",
-  model,
-  provider,
-  configs,
+  model: defaultModel,
+  provider: defaultProvider,
+  configs: defaultConfigs,
   workspaceName,
   onAccept,
 }) => {
@@ -79,6 +83,14 @@ const PromptImprovementDialog: React.FC<PromptImprovementDialogProps> = ({
   const [isEditorFocused, setIsEditorFocused] = useState(false);
   const editorViewRef = useRef<EditorView | null>(null);
 
+  // Model selection with persistence using the reusable hook
+  const { model, provider, configs, modelSelectProps } = useModelSelection({
+    persistenceKey: PROMPT_IMPROVEMENT_LAST_PICKED_MODEL,
+    defaultModel,
+    defaultProvider,
+    defaultConfigs,
+  });
+
   const { improvePrompt, generatePrompt } = usePromptImprovement({
     workspaceName,
   });
@@ -443,6 +455,16 @@ const PromptImprovementDialog: React.FC<PromptImprovementDialogProps> = ({
           />
         </DialogHeader>
         <DialogAutoScrollBody>
+          <div className="mb-4 flex items-center gap-3">
+            <div className="comet-body-accented shrink-0">Model</div>
+            <div className="w-64">
+              <PromptModelSelect
+                workspaceName={workspaceName}
+                {...modelSelectProps}
+                disabled={isLoading}
+              />
+            </div>
+          </div>
           {error && (
             <Alert variant="destructive" className="mb-4">
               <AlertTitle>{error}</AlertTitle>
diff --git a/apps/opik-frontend/src/hooks/useModelSelection.ts b/apps/opik-frontend/src/hooks/useModelSelection.ts
@@ -0,0 +1,163 @@
+import { useCallback, useMemo } from "react";
+import useLastPickedModel from "@/hooks/useLastPickedModel";
+import useLLMProviderModelsData from "@/hooks/useLLMProviderModelsData";
+import useProviderKeys from "@/api/provider-keys/useProviderKeys";
+import useAppStore from "@/store/AppStore";
+import { getDefaultConfigByProvider } from "@/lib/playground";
+import {
+  COMPOSED_PROVIDER_TYPE,
+  LLMPromptConfigsType,
+  PROVIDER_MODEL_TYPE,
+} from "@/types/providers";
+
+export interface UseModelSelectionParams {
+  persistenceKey: string;
+  defaultModel?: string;
+  defaultProvider?: COMPOSED_PROVIDER_TYPE | "";
+  defaultConfigs?: LLMPromptConfigsType;
+}
+
+export interface ModelSelectProps {
+  value: PROVIDER_MODEL_TYPE | "";
+  provider: COMPOSED_PROVIDER_TYPE | "";
+  onChange: (
+    model: PROVIDER_MODEL_TYPE,
+    provider: COMPOSED_PROVIDER_TYPE,
+  ) => void;
+  onAddProvider: (provider: COMPOSED_PROVIDER_TYPE) => void;
+  onDeleteProvider: (provider: COMPOSED_PROVIDER_TYPE) => void;
+}
+
+export interface UseModelSelectionResult {
+  model: PROVIDER_MODEL_TYPE | "";
+  provider: COMPOSED_PROVIDER_TYPE | "";
+  configs: LLMPromptConfigsType;
+  modelSelectProps: ModelSelectProps;
+}
+
+const useModelSelection = ({
+  persistenceKey,
+  defaultModel,
+  defaultProvider,
+  defaultConfigs,
+}: UseModelSelectionParams): UseModelSelectionResult => {
+  const workspaceName = useAppStore((state) => state.activeWorkspaceName);
+
+  const [lastPickedModel, setLastPickedModel] = useLastPickedModel({
+    key: persistenceKey,
+  });
+
+  const { data: providerKeysData } = useProviderKeys({
+    workspaceName,
+  });
+
+  const providerKeys = useMemo(() => {
+    return providerKeysData?.content?.map((c) => c.ui_composed_provider) || [];
+  }, [providerKeysData]);
+
+  const { calculateModelProvider, calculateDefaultModel } =
+    useLLMProviderModelsData();
+
+  const { model, provider, configs } = useMemo(() => {
+    if (lastPickedModel) {
+      const lastPickedProvider = calculateModelProvider(lastPickedModel);
+      if (lastPickedProvider && providerKeys.includes(lastPickedProvider)) {
+        return {
+          model: lastPickedModel,
+          provider: lastPickedProvider,
+          configs: getDefaultConfigByProvider(
+            lastPickedProvider,
+            lastPickedModel,
+          ),
+        };
+      }
+    }
+
+    if (defaultModel && defaultProvider) {
+      return {
+        model: defaultModel as PROVIDER_MODEL_TYPE | "",
+        provider: defaultProvider,
+        configs: defaultConfigs ?? getDefaultConfigByProvider(defaultProvider),
+      };
+    }
+
+    const calculatedModel = calculateDefaultModel(
+      lastPickedModel,
+      providerKeys,
+    ) as PROVIDER_MODEL_TYPE | "";
+    const calculatedProvider = calculateModelProvider(calculatedModel);
+    return {
+      model: calculatedModel,
+      provider: calculatedProvider,
+      configs: getDefaultConfigByProvider(calculatedProvider, calculatedModel),
+    };
+  }, [
+    lastPickedModel,
+    providerKeys,
+    calculateModelProvider,
+    calculateDefaultModel,
+    defaultModel,
+    defaultProvider,
+    defaultConfigs,
+  ]);
+
+  const handleModelChange = useCallback(
+    (newModel: PROVIDER_MODEL_TYPE) => {
+      setLastPickedModel(newModel);
+    },
+    [setLastPickedModel],
+  );
+
+  const handleAddProvider = useCallback(
+    (addedProvider: COMPOSED_PROVIDER_TYPE) => {
+      if (!model) {
+        setLastPickedModel(
+          calculateDefaultModel(
+            model as PROVIDER_MODEL_TYPE | "",
+            [addedProvider],
+            addedProvider,
+          ),
+        );
+      }
+    },
+    [calculateDefaultModel, model, setLastPickedModel],
+  );
+
+  const handleDeleteProvider = useCallback(
+    (deletedProvider: COMPOSED_PROVIDER_TYPE) => {
+      const currentProvider = calculateModelProvider(
+        model as PROVIDER_MODEL_TYPE | "",
+      );
+      if (currentProvider === deletedProvider) {
+        setLastPickedModel("");
+      }
+    },
+    [calculateModelProvider, model, setLastPickedModel],
+  );
+
+  const modelSelectProps: ModelSelectProps = useMemo(
+    () => ({
+      value: model as PROVIDER_MODEL_TYPE | "",
+      provider,
+      onChange: handleModelChange,
+      onAddProvider: handleAddProvider,
+      onDeleteProvider: handleDeleteProvider,
+    }),
+    [
+      model,
+      provider,
+      handleModelChange,
+      handleAddProvider,
+      handleDeleteProvider,
+    ],
+  );
+
+  return {
+    model: model as PROVIDER_MODEL_TYPE | "",
+    provider,
+    configs,
+    modelSelectProps,
+  };
+};
+
+export default useModelSelection;
