fallback to other saved providers if default is not there

Arize-ai · Oct 29, 2024 · 8a08447 · 8a08447
1 parent efe0731
commit 8a08447
Show file tree

Hide file tree

Showing 3 changed files with 128 additions and 9 deletions.
diff --git a/app/src/constants/generativeConstants.ts b/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";
 

diff --git a/app/src/store/playground/__tests__/playgroundStore.test.ts b/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");
+  });
+});
diff --git a/app/src/store/playground/playgroundStore.tsx b/app/src/store/playground/playgroundStore.tsx
@@ -5,6 +5,7 @@ import { TemplateLanguages } from "@phoenix/components/templateEditor/constants"
 import { TemplateLanguage } from "@phoenix/components/templateEditor/types";
 import {
   DEFAULT_CHAT_ROLE,
+  DEFAULT_MODEL_NAME,
   DEFAULT_MODEL_PROVIDER,
 } from "@phoenix/constants/generativeConstants";
 import { OpenAIToolCall } from "@phoenix/schemas";
@@ -96,7 +97,7 @@ export function createPlaygroundInstance(): PlaygroundInstance {
     template: generateChatCompletionTemplate(),
     model: {
       provider: DEFAULT_MODEL_PROVIDER,
-      modelName: "gpt-4o",
+      modelName: DEFAULT_MODEL_NAME,
       invocationParameters: {},
     },
     tools: [],
@@ -155,20 +156,28 @@ export function createOpenAITool(toolNumber: number): OpenAITool {
  * If the initial props has instances, those will be used
  * If not a single instance will be created and saved model config defaults will be used
  * @returns a list of {@link PlaygroundInstance} instances
+ *
+ * NB: This function is only exported for testing
  */
-function getInitialInstances(initialProps: InitialPlaygroundState) {
+export function getInitialInstances(initialProps: InitialPlaygroundState) {
   if (initialProps.instances != null && initialProps.instances.length > 0) {
     return initialProps.instances;
   }
   const instance = createPlaygroundInstance();
-  const savedDefaultProviderConfig =
-    initialProps.modelConfigByProvider[instance.model.provider];
-  if (savedDefaultProviderConfig) {
-    instance.model = {
-      ...instance.model,
-      ...savedDefaultProviderConfig,
-    };
+
+  const savedModelConfigs = Object.values(initialProps.modelConfigByProvider);
+  const hasSavedModelConfig = savedModelConfigs.length > 0;
+  if (!hasSavedModelConfig) {
+    return [instance];
   }
+  const savedDefaultProviderConfig =
+    savedModelConfigs.find(
+      (config) => config.provider === DEFAULT_MODEL_PROVIDER
+    ) ?? savedModelConfigs[0];
+  instance.model = {
+    ...instance.model,
+    ...savedDefaultProviderConfig,
+  };
   return [instance];
 }