feat(playground): save model config by provider in preferences

app/src/constants/generativeConstants.ts

@@ -11,6 +11,10 @@ export const ModelProviders: Record<ModelProvider, string> = {
  * The default model provider
  */
 export const DEFAULT_MODEL_PROVIDER: ModelProvider = "OPENAI";
+/**
+ * The default model name
+ */
+export const DEFAULT_MODEL_NAME = "gpt-4o";
 
 export const DEFAULT_CHAT_ROLE: ChatMessageRole = "user";
 

app/src/contexts/PlaygroundContext.tsx

@@ -3,7 +3,7 @@ import { useZustand } from "use-zustand";
 
 import {
   createPlaygroundStore,
-  PlaygroundProps,
+  InitialPlaygroundState,
   PlaygroundState,
   PlaygroundStore,
 } from "@phoenix/store";
@@ -13,7 +13,7 @@ export const PlaygroundContext = createContext<PlaygroundStore | null>(null);
 export function PlaygroundProvider({
   children,
   ...props
-}: PropsWithChildren<Partial<PlaygroundProps>>) {
+}: PropsWithChildren<InitialPlaygroundState>) {
   const [store] = useState<PlaygroundStore>(() => createPlaygroundStore(props));
   return (
     <PlaygroundContext.Provider value={store}>

app/src/pages/playground/ModelConfigButton.tsx

@@ -18,14 +18,18 @@ import {
   Picker,
   Text,
   TextField,
+  Tooltip,
+  TooltipTrigger,
   View,
 } from "@arizeai/components";
 
 import {
   AZURE_OPENAI_API_VERSIONS,
   ModelProviders,
 } from "@phoenix/constants/generativeConstants";
+import { useNotifySuccess } from "@phoenix/contexts";
 import { usePlaygroundContext } from "@phoenix/contexts/PlaygroundContext";
+import { usePreferencesContext } from "@phoenix/contexts/PreferencesContext";
 import { PlaygroundInstance } from "@phoenix/store";
 
 import { ModelConfigButtonDialogQuery } from "./__generated__/ModelConfigButtonDialogQuery.graphql";
@@ -43,6 +47,9 @@ function AzureOpenAiModelConfigFormField({
   instance: PlaygroundInstance;
 }) {
   const updateModel = usePlaygroundContext((state) => state.updateModel);
+  const modelConfigByProvider = usePreferencesContext(
+    (state) => state.modelConfigByProvider
+  );
 
   const updateModelConfig = useCallback(
     ({
@@ -58,9 +65,10 @@ function AzureOpenAiModelConfigFormField({
           ...instance.model,
           [configKey]: value,
         },
+        modelConfigByProvider,
       });
     },
-    [instance.id, instance.model, updateModel]
+    [instance.id, instance.model, modelConfigByProvider, updateModel]
   );
 
   return (
@@ -128,13 +136,7 @@ export function ModelConfigButton(props: ModelConfigButtonProps) {
         size="compact"
         onClick={() => {
           startTransition(() => {
-            setDialog(
-              <Dialog title="Model Configuration" size="M">
-                <Suspense>
-                  <ModelConfigDialogContent {...props} />
-                </Suspense>
-              </Dialog>
-            );
+            setDialog(<ModelConfigDialog {...props} />);
           });
         }}
       >
@@ -156,18 +158,75 @@ export function ModelConfigButton(props: ModelConfigButtonProps) {
   );
 }
 
+interface ModelConfigDialogProps extends ModelConfigButtonProps {}
+function ModelConfigDialog(props: ModelConfigDialogProps) {
+  const instance = usePlaygroundContext((state) =>
+    state.instances.find(
+      (instance) => instance.id === props.playgroundInstanceId
+    )
+  );
+
+  if (!instance) {
+    throw new Error(
+      `Playground instance ${props.playgroundInstanceId} not found`
+    );
+  }
+  const setModelConfigForProvider = usePreferencesContext(
+    (state) => state.setModelConfigForProvider
+  );
+
+  const notifySuccess = useNotifySuccess();
+  const onSaveConfig = useCallback(() => {
+    setModelConfigForProvider({
+      provider: instance.model.provider,
+      modelConfig: instance.model,
+    });
+    notifySuccess({
+      title: "Model Configuration Saved",
+      message: `${ModelProviders[instance.model.provider]} model configuration saved`,
+      expireMs: 3000,
+    });
+  }, [instance.model, notifySuccess, setModelConfigForProvider]);
+  return (
+    <Dialog
+      title="Model Configuration"
+      size="M"
+      extra={
+        <TooltipTrigger delay={0} offset={5}>
+          <Button size={"compact"} variant="default" onClick={onSaveConfig}>
+            Save Config
+          </Button>
+          <Tooltip>
+            Remember configuration for{" "}
+            {ModelProviders[instance.model.provider] ?? "this provider"}.
+          </Tooltip>
+        </TooltipTrigger>
+      }
+    >
+      <Suspense>
+        <ModelConfigDialogContent {...props} />
+      </Suspense>
+    </Dialog>
+  );
+}
+
 interface ModelConfigDialogContentProps extends ModelConfigButtonProps {}
 function ModelConfigDialogContent(props: ModelConfigDialogContentProps) {
   const { playgroundInstanceId } = props;
-  const updateModel = usePlaygroundContext((state) => state.updateModel);
   const instance = usePlaygroundContext((state) =>
     state.instances.find((instance) => instance.id === playgroundInstanceId)
   );
+
   if (!instance) {
     throw new Error(
       `Playground instance ${props.playgroundInstanceId} not found`
     );
   }
+  const modelConfigByProvider = usePreferencesContext(
+    (state) => state.modelConfigByProvider
+  );
+  const updateModel = usePlaygroundContext((state) => state.updateModel);
+
   const query = useLazyLoadQuery<ModelConfigButtonDialogQuery>(
     graphql`
       query ModelConfigButtonDialogQuery($providerKey: GenerativeProviderKey!) {
@@ -186,9 +245,15 @@ function ModelConfigDialogContent(props: ModelConfigDialogContentProps) {
           provider: instance.model.provider,
           modelName,
         },
+        modelConfigByProvider,
       });
     },
-    [instance.model.provider, playgroundInstanceId, updateModel]
+    [
+      instance.model.provider,
+      modelConfigByProvider,
+      playgroundInstanceId,
+      updateModel,
+    ]
   );
 
   const onInvocationParametersChange: InvocationParametersChangeHandler =
@@ -203,9 +268,10 @@ function ModelConfigDialogContent(props: ModelConfigDialogContentProps) {
               [parameter]: value,
             },
           },
+          modelConfigByProvider,
         });
       },
-      [instance.model, playgroundInstanceId, updateModel]
+      [instance.model, modelConfigByProvider, playgroundInstanceId, updateModel]
     );
 
   return (
@@ -221,6 +287,7 @@ function ModelConfigDialogContent(props: ModelConfigDialogContentProps) {
                 provider,
                 modelName: null,
               },
+              modelConfigByProvider,
             });
           }}
         />

app/src/pages/playground/Playground.tsx

@@ -20,7 +20,8 @@ import {
   PlaygroundProvider,
   usePlaygroundContext,
 } from "@phoenix/contexts/PlaygroundContext";
-import { InitialPlaygroundState } from "@phoenix/store";
+import { usePreferencesContext } from "@phoenix/contexts/PreferencesContext";
+import { PlaygroundProps } from "@phoenix/store";
 
 import { NUM_MAX_PLAYGROUND_INSTANCES } from "./constants";
 import { PlaygroundCredentialsDropdown } from "./PlaygroundCredentialsDropdown";
@@ -39,7 +40,10 @@ const playgroundWrapCSS = css`
   height: 100%;
 `;
 
-export function Playground(props: InitialPlaygroundState) {
+export function Playground(props: Partial<PlaygroundProps>) {
+  const modelConfigByProvider = usePreferencesContext(
+    (state) => state.modelConfigByProvider
+  );
   const showStreamToggle = useFeatureFlag("playgroundNonStreaming");
   const [, setSearchParams] = useSearchParams();
 
@@ -56,7 +60,10 @@ export function Playground(props: InitialPlaygroundState) {
   }, [setSearchParams]);
 
   return (
-    <PlaygroundProvider {...props}>
+    <PlaygroundProvider
+      {...props}
+      modelConfigByProvider={modelConfigByProvider}
+    >
       <div css={playgroundWrapCSS}>
         <View
           borderBottomColor="dark"

app/src/store/playground/__tests__/playgroundStore.test.ts

@@ -0,0 +1,106 @@
+import {
+  DEFAULT_MODEL_NAME,
+  DEFAULT_MODEL_PROVIDER,
+} from "@phoenix/constants/generativeConstants";
+
+import {
+  _resetInstanceId,
+  createPlaygroundInstance,
+  getInitialInstances,
+} from "../playgroundStore";
+import { InitialPlaygroundState } from "../types";
+
+describe("getInitialInstances", () => {
+  beforeEach(() => {
+    _resetInstanceId();
+  });
+  it("should return instances from initialProps if they exist", () => {
+    const existingInstance = {
+      ...createPlaygroundInstance(),
+      model: {
+        modelName: "test-model",
+        provider: "OPENAI" as const,
+        invocationParameters: {},
+      },
+    };
+    const initialProps: InitialPlaygroundState = {
+      instances: [existingInstance],
+      modelConfigByProvider: {},
+    };
+
+    const instances = getInitialInstances(initialProps);
+
+    expect(instances).toEqual([existingInstance]);
+  });
+
+  it("should create a new default instance if no instances exist in initialProps and there are no saved modelConfigs", () => {
+    const initialProps: InitialPlaygroundState = {
+      modelConfigByProvider: {},
+    };
+    const instances = getInitialInstances(initialProps);
+
+    expect(instances).toHaveLength(1);
+    expect(instances[0].id).toBe(0);
+    expect(instances[0].model.provider).toBe(DEFAULT_MODEL_PROVIDER);
+    expect(instances[0].model.modelName).toBe(DEFAULT_MODEL_NAME);
+  });
+
+  it("should use saved model config if available", () => {
+    const initialProps: InitialPlaygroundState = {
+      modelConfigByProvider: {
+        OPENAI: {
+          modelName: "test-model",
+          provider: "OPENAI",
+          invocationParameters: {},
+        },
+      },
+    };
+
+    const instances = getInitialInstances(initialProps);
+
+    expect(instances).toHaveLength(1);
+    expect(instances[0].model.provider).toBe("OPENAI");
+    expect(instances[0].model.modelName).toBe("test-model");
+  });
+
+  it("should use default model provider config if available", () => {
+    const initialProps: InitialPlaygroundState = {
+      modelConfigByProvider: {
+        OPENAI: {
+          modelName: "test-model-openai",
+          provider: "OPENAI",
+          invocationParameters: {},
+        },
+        ANTHROPIC: {
+          modelName: "test-model-anthropic",
+          provider: "ANTHROPIC",
+          invocationParameters: {},
+        },
+      },
+    };
+
+    const instances = getInitialInstances(initialProps);
+
+    expect(instances).toHaveLength(1);
+    expect(instances[0].model.provider).toBe("OPENAI");
+    expect(instances[0].model.modelName).toBe("test-model-openai");
+  });
+
+  it("should use any saved config if available if the default provider config is not", () => {
+    const initialProps: InitialPlaygroundState = {
+      modelConfigByProvider: {
+        ANTHROPIC: {
+          modelName: "test-model-anthropic",
+          provider: "ANTHROPIC",
+          invocationParameters: {},
+        },
+      },
+    };
+
+    const instances = getInitialInstances(initialProps);
+
+    expect(instances).toHaveLength(1);
+    expect(instances[0].model.provider).toBe("ANTHROPIC");
+    expect(instances[0].model.modelName).toBe("test-model-anthropic");
+  });
+});