From 6fc0ffe3ffabfdb2a7d4e5d66bddc83ee1eaf272 Mon Sep 17 00:00:00 2001
From: Sasha <aliaksandr@comet.com>
Date: Mon, 13 Jan 2025 13:12:37 +0100
Subject: [PATCH 1/2] [OPIK-740]: handle adding an ai provider in the
 playground;

---
 .../PlaygroundPrompts/PlaygroundPrompt.tsx    | 22 +++++++++++++++++++
 .../PromptModelSelect/PromptModelSelect.tsx   |  3 +++
 .../AddEditAIProviderDialog.tsx               | 17 +++++++++++++-
 3 files changed, 41 insertions(+), 1 deletion(-)

diff --git a/apps/opik-frontend/src/components/pages/PlaygroundPage/PlaygroundPrompts/PlaygroundPrompt.tsx b/apps/opik-frontend/src/components/pages/PlaygroundPage/PlaygroundPrompts/PlaygroundPrompt.tsx
index 8c34ba421c..d1178edb57 100644
--- a/apps/opik-frontend/src/components/pages/PlaygroundPage/PlaygroundPrompts/PlaygroundPrompt.tsx
+++ b/apps/opik-frontend/src/components/pages/PlaygroundPage/PlaygroundPrompts/PlaygroundPrompt.tsx
@@ -118,6 +118,27 @@ const PlaygroundPrompt = ({
     [updatePrompt, promptId],
   );
 
+  const handleAddProvider = useCallback(
+    (provider: PROVIDER_TYPE) => {
+      const modelProvider = model ? getModelProvider(model) : "";
+      const noCurrentModel = !modelProvider;
+
+      if (noCurrentModel) {
+        const newModel = PROVIDERS[provider].defaultModel;
+
+        const newDefaultConfigs = provider
+          ? getDefaultConfigByProvider(provider)
+          : {};
+
+        updatePrompt(promptId, {
+          model: newModel,
+          configs: newDefaultConfigs,
+        });
+      }
+    },
+    [model],
+  );
+
   useEffect(() => {
     // on init, to check if a prompt has a model from valid providers: (f.e., remove a provider after setting a model)
     if (!checkedIfModelIsValidRef.current && !isPendingProviderKeys) {
@@ -175,6 +196,7 @@ const PlaygroundPrompt = ({
               onChange={handleUpdateModel}
               provider={provider}
               workspaceName={workspaceName}
+              onAddProvider={handleAddProvider}
             />
           </div>
           <PromptModelConfigs
diff --git a/apps/opik-frontend/src/components/pages/PlaygroundPage/PlaygroundPrompts/PromptModelSelect/PromptModelSelect.tsx b/apps/opik-frontend/src/components/pages/PlaygroundPage/PlaygroundPrompts/PromptModelSelect/PromptModelSelect.tsx
index af5c4de087..9fc3443bae 100644
--- a/apps/opik-frontend/src/components/pages/PlaygroundPage/PlaygroundPrompts/PromptModelSelect/PromptModelSelect.tsx
+++ b/apps/opik-frontend/src/components/pages/PlaygroundPage/PlaygroundPrompts/PromptModelSelect/PromptModelSelect.tsx
@@ -35,6 +35,7 @@ interface PromptModelSelectProps {
   workspaceName: string;
   onChange: (value: PROVIDER_MODEL_TYPE) => void;
   provider: PROVIDER_TYPE | "";
+  onAddProvider?: (provider: PROVIDER_TYPE) => void;
 }
 
 const STALE_TIME = 1000;
@@ -44,6 +45,7 @@ const PromptModelSelect = ({
   workspaceName,
   onChange,
   provider,
+  onAddProvider,
 }: PromptModelSelectProps) => {
   const resetDialogKeyRef = useRef(0);
   const inputRef = useRef<HTMLInputElement>(null);
@@ -290,6 +292,7 @@ const PromptModelSelect = ({
         key={resetDialogKeyRef.current}
         open={openConfigDialog}
         setOpen={setOpenConfigDialog}
+        onAddProvider={onAddProvider}
       />
     </>
   );
diff --git a/apps/opik-frontend/src/components/shared/AddEditAIProviderDialog/AddEditAIProviderDialog.tsx b/apps/opik-frontend/src/components/shared/AddEditAIProviderDialog/AddEditAIProviderDialog.tsx
index 921e023f6d..d54d81a305 100644
--- a/apps/opik-frontend/src/components/shared/AddEditAIProviderDialog/AddEditAIProviderDialog.tsx
+++ b/apps/opik-frontend/src/components/shared/AddEditAIProviderDialog/AddEditAIProviderDialog.tsx
@@ -17,17 +17,20 @@ import { PROVIDERS, PROVIDERS_OPTIONS } from "@/constants/providers";
 import { SelectItem } from "@/components/ui/select";
 import { DropdownOption } from "@/types/shared";
 import EyeInput from "@/components/shared/EyeInput/EyeInput";
+import isFunction from "lodash/isFunction";
 
 type AddEditAIProviderDialogProps = {
   providerKey?: ProviderKey;
   open: boolean;
   setOpen: (open: boolean) => void;
+  onAddProvider?: (provider: PROVIDER_TYPE) => void;
 };
 
 const AddEditAIProviderDialog: React.FC<AddEditAIProviderDialogProps> = ({
   providerKey,
   open,
   setOpen,
+  onAddProvider,
 }) => {
   const { mutate: createMutate } = useProviderKeysCreateMutation();
   const { mutate: updateMutate } = useProviderKeysUpdateMutation();
@@ -58,6 +61,10 @@ const AddEditAIProviderDialog: React.FC<AddEditAIProviderDialogProps> = ({
         },
       });
     } else if (provider) {
+      if (isFunction(onAddProvider)) {
+        onAddProvider(provider);
+      }
+
       createMutate({
         providerKey: {
           apiKey,
@@ -65,7 +72,15 @@ const AddEditAIProviderDialog: React.FC<AddEditAIProviderDialogProps> = ({
         },
       });
     }
-  }, [createMutate, isEdit, apiKey, updateMutate, provider, providerKey]);
+  }, [
+    createMutate,
+    isEdit,
+    apiKey,
+    updateMutate,
+    provider,
+    providerKey,
+    onAddProvider,
+  ]);
 
   const renderOption = (option: DropdownOption<string>) => {
     const Icon = PROVIDERS[option.value as PROVIDER_TYPE].icon;

From 43de1cb4d175eac50acd1b7ae8b13189a0405821 Mon Sep 17 00:00:00 2001
From: Sasha <aliaksandr@comet.com>
Date: Mon, 13 Jan 2025 13:13:26 +0100
Subject: [PATCH 2/2] [OPIK-740]: add dependencies to handleAddProvider;

---
 .../pages/PlaygroundPage/PlaygroundPrompts/PlaygroundPrompt.tsx | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/apps/opik-frontend/src/components/pages/PlaygroundPage/PlaygroundPrompts/PlaygroundPrompt.tsx b/apps/opik-frontend/src/components/pages/PlaygroundPage/PlaygroundPrompts/PlaygroundPrompt.tsx
index d1178edb57..6b02048f01 100644
--- a/apps/opik-frontend/src/components/pages/PlaygroundPage/PlaygroundPrompts/PlaygroundPrompt.tsx
+++ b/apps/opik-frontend/src/components/pages/PlaygroundPage/PlaygroundPrompts/PlaygroundPrompt.tsx
@@ -136,7 +136,7 @@ const PlaygroundPrompt = ({
         });
       }
     },
-    [model],
+    [model, promptId, updatePrompt],
   );
 
   useEffect(() => {