[OPIK-740]: handle adding provider (#1031)

* [OPIK-740]: handle adding an ai provider in the playground; * [OPIK-740]: add dependencies to handleAddProvider; --------- Co-authored-by: Sasha <aliaksandr@comet.com>
comet-ml · Jan 13, 2025 · 4489475 · 4489475
1 parent c09dc3a
commit 4489475
Show file tree

Hide file tree

Showing 3 changed files with 41 additions and 1 deletion.
diff --git a/.../opik-frontend/src/components/pages/PlaygroundPage/PlaygroundPrompts/PlaygroundPrompt.tsx b/.../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, promptId, updatePrompt],
+  );
+
   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/...components/pages/PlaygroundPage/PlaygroundPrompts/PromptModelSelect/PromptModelSelect.tsx b/...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
@@ -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,14 +61,26 @@ const AddEditAIProviderDialog: React.FC<AddEditAIProviderDialogProps> = ({
         },
       });
     } else if (provider) {
+      if (isFunction(onAddProvider)) {
+        onAddProvider(provider);
+      }
+
       createMutate({
         providerKey: {
           apiKey,
           provider,
         },
       });
     }
-  }, [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;