feat(playground): add credential storage (#4970)

* feat(playground): add credential storage

* organize

* reorganize store files

* add credentials dropdown

* remove unused icon

* make credentials inputs unique per provider

* update unique logic for providers

* migrate to instance.model.provider from instance.provider

* move provider in tests

app/src/@types/generative.d.ts

@@ -1 +1,6 @@
 declare type ModelProvider = "OPENAI" | "AZURE_OPENAI" | "ANTHROPIC";
+
+/**
+ * The role of a chat message
+ */
+declare type ChatMessageRole = "user" | "system" | "ai" | "tool";

app/src/App.tsx

@@ -4,6 +4,7 @@ import { ThemeProvider as EmotionThemeProvider } from "@emotion/react";
 
 import { Provider, theme } from "@arizeai/components";
 
+import { CredentialsProvider } from "./contexts/CredentialsContext";
 import { FeatureFlagsProvider } from "./contexts/FeatureFlagsContext";
 import { FunctionalityProvider } from "./contexts/FunctionalityContext";
 import { PreferencesProvider } from "./contexts/PreferencesContext";
@@ -33,11 +34,13 @@ export function AppContent() {
           <GlobalStyles />
           <FeatureFlagsProvider>
             <PreferencesProvider>
-              <Suspense>
-                <NotificationProvider>
-                  <AppRoutes />
-                </NotificationProvider>
-              </Suspense>
+              <CredentialsProvider>
+                <Suspense>
+                  <NotificationProvider>
+                    <AppRoutes />
+                  </NotificationProvider>
+                </Suspense>
+              </CredentialsProvider>
             </PreferencesProvider>
           </FeatureFlagsProvider>
         </RelayEnvironmentProvider>

app/src/contexts/CredentialsContext.tsx

@@ -0,0 +1,35 @@
+import React, { createContext, PropsWithChildren, useState } from "react";
+import { useZustand } from "use-zustand";
+
+import {
+  createCredentialsStore,
+  CredentialsProps,
+  CredentialsState,
+  CredentialsStore,
+} from "@phoenix/store";
+
+export const CredentialsContext = createContext<CredentialsStore | null>(null);
+
+export function CredentialsProvider({
+  children,
+  ...props
+}: PropsWithChildren<Partial<CredentialsProps>>) {
+  const [store] = useState<CredentialsStore>(() =>
+    createCredentialsStore(props)
+  );
+  return (
+    <CredentialsContext.Provider value={store}>
+      {children}
+    </CredentialsContext.Provider>
+  );
+}
+
+export function useCredentialsContext<T>(
+  selector: (state: CredentialsState) => T,
+  equalityFn?: (left: T, right: T) => boolean
+): T {
+  const store = React.useContext(CredentialsContext);
+  if (!store)
+    throw new Error("Missing CredentialsContext.Provider in the tree");
+  return useZustand(store, selector, equalityFn);
+}

app/src/contexts/PlaygroundContext.tsx

@@ -6,7 +6,7 @@ import {
   PlaygroundProps,
   PlaygroundState,
   PlaygroundStore,
-} from "@phoenix/store/playgroundStore";
+} from "@phoenix/store";
 
 export const PlaygroundContext = createContext<PlaygroundStore | null>(null);
 

app/src/pages/playground/MessageRolePicker.tsx

@@ -3,8 +3,6 @@ import { css } from "@emotion/react";
 
 import { Item, Picker } from "@arizeai/components";
 
-import { ChatMessageRole } from "@phoenix/store";
-
 import { isChatMessageRole } from "./playgroundUtils";
 
 const hiddenLabelCSS = css`

app/src/pages/playground/Playground.tsx

@@ -11,6 +11,7 @@ import {
 import { InitialPlaygroundState } from "@phoenix/store";
 
 import { NUM_MAX_PLAYGROUND_INSTANCES } from "./constants";
+import { PlaygroundCredentialsDropdown } from "./PlaygroundCredentialsDropdown";
 import { PlaygroundInputTypeTypeRadioGroup } from "./PlaygroundInputModeRadioGroup";
 import { PlaygroundInstance } from "./PlaygroundInstance";
 import { PlaygroundRunButton } from "./PlaygroundRunButton";
@@ -33,9 +34,7 @@ export function Playground(props: InitialPlaygroundState) {
             <Heading level={1}>Playground</Heading>
             <Flex direction="row" gap="size-100" alignItems="center">
               <PlaygroundInputTypeTypeRadioGroup />
-              <Button variant="default" size="compact">
-                API Keys
-              </Button>
+              <PlaygroundCredentialsDropdown />
               <AddPromptButton />
               <PlaygroundRunButton />
             </Flex>

app/src/pages/playground/PlaygroundChatTemplate.tsx

@@ -26,7 +26,6 @@ import { usePlaygroundContext } from "@phoenix/contexts/PlaygroundContext";
 import { useChatMessageStyles } from "@phoenix/hooks/useChatMessageStyles";
 import {
   ChatMessage,
-  ChatMessageRole,
   generateMessageId,
   PlaygroundChatTemplate as PlaygroundChatTemplateType,
 } from "@phoenix/store";
@@ -140,7 +139,7 @@ export function PlaygroundChatTemplate(props: PlaygroundChatTemplateProps) {
                       ...template.messages,
                       {
                         id: generateMessageId(),
-                        role: ChatMessageRole.user,
+                        role: "user",
                         content: "",
                       },
                     ],

app/src/pages/playground/PlaygroundCredentialsDropdown.tsx

@@ -0,0 +1,69 @@
+import React from "react";
+
+import {
+  DropdownButton,
+  DropdownMenu,
+  DropdownTrigger,
+  Flex,
+  Form,
+  Heading,
+  Text,
+  TextField,
+  View,
+} from "@arizeai/components";
+
+import { useCredentialsContext } from "@phoenix/contexts/CredentialsContext";
+import { usePlaygroundContext } from "@phoenix/contexts/PlaygroundContext";
+import { CredentialKey } from "@phoenix/store";
+
+export const ProviderToCredentialKeyMap: Record<ModelProvider, CredentialKey> =
+  {
+    OPENAI: "OPENAI_API_KEY",
+    ANTHROPIC: "ANTHROPIC_API_KEY",
+    AZURE_OPENAI: "AZURE_OPENAI_API_KEY",
+  };
+
+export function PlaygroundCredentialsDropdown() {
+  const currentProviders = usePlaygroundContext((state) =>
+    Array.from(
+      new Set(state.instances.map((instance) => instance.model.provider))
+    )
+  );
+  const setCredential = useCredentialsContext((state) => state.setCredential);
+  const credentials = useCredentialsContext((state) => state);
+  return (
+    <DropdownTrigger placement="bottom">
+      <DropdownButton size="compact">API Keys</DropdownButton>
+      <DropdownMenu>
+        <View padding="size-200">
+          <Flex direction={"column"} gap={"size-100"}>
+            <Heading level={2} weight="heavy">
+              API Keys
+            </Heading>
+            <Text color="white70">
+              API keys are stored in your browser and used to communicate with
+              their respective API&apos;s.
+            </Text>
+            <Form>
+              {currentProviders.map((provider) => {
+                const credentialKey = ProviderToCredentialKeyMap[provider];
+                return (
+                  <TextField
+                    key={provider}
+                    label={credentialKey}
+                    type="password"
+                    isRequired
+                    onChange={(value) => {
+                      setCredential({ credential: credentialKey, value });
+                    }}
+                    value={credentials[credentialKey]}
+                  />
+                );
+              })}
+            </Form>
+          </Flex>
+        </View>
+      </DropdownMenu>
+    </DropdownTrigger>
+  );
+}

app/src/pages/playground/PlaygroundOutput.tsx

@@ -6,11 +6,7 @@ import { Card, Flex, Icon, Icons } from "@arizeai/components";
 
 import { usePlaygroundContext } from "@phoenix/contexts/PlaygroundContext";
 import { useChatMessageStyles } from "@phoenix/hooks/useChatMessageStyles";
-import {
-  ChatMessage,
-  ChatMessageRole,
-  generateMessageId,
-} from "@phoenix/store";
+import { ChatMessage, generateMessageId } from "@phoenix/store";
 import { assertUnreachable } from "@phoenix/typeUtils";
 
 import {
@@ -70,7 +66,7 @@ export function PlaygroundOutput(props: PlaygroundOutputProps) {
           message={{
             id: generateMessageId(),
             content: instance.output,
-            role: ChatMessageRole.ai,
+            role: "ai",
           }}
         />
       );
@@ -144,13 +140,13 @@ function toGqlChatCompletionRole(
   role: ChatMessageRole
 ): ChatCompletionMessageRole {
   switch (role) {
-    case ChatMessageRole.system:
+    case "system":
       return "SYSTEM";
-    case ChatMessageRole.user:
+    case "user":
       return "USER";
-    case ChatMessageRole.tool:
+    case "tool":
       return "TOOL";
-    case ChatMessageRole.ai:
+    case "ai":
       return "AI";
     default:
       assertUnreachable(role);
@@ -204,7 +200,7 @@ function PlaygroundOutputText(props: PlaygroundInstanceProps) {
       message={{
         id: generateMessageId(),
         content: output,
-        role: ChatMessageRole.ai,
+        role: "ai",
       }}
     />
   );

app/src/pages/playground/__tests__/playgroundUtils.test.ts

@@ -1,7 +1,6 @@
 import {
   _resetInstanceId,
   _resetMessageId,
-  ChatMessageRole,
   PlaygroundInstance,
 } from "@phoenix/store";
 
@@ -36,26 +35,24 @@ const expectedPlaygroundInstanceWithIO: PlaygroundInstance = {
     // These id's are not 0, 1, 2, because we create a playground instance (including messages) at the top of the transformSpanAttributesToPlaygroundInstance function
     // Doing so increments the message id counter
     messages: [
-      { id: 2, content: "You are a chatbot", role: ChatMessageRole.system },
-      { id: 3, content: "hello?", role: ChatMessageRole.user },
+      { id: 2, content: "You are a chatbot", role: "system" },
+      { id: 3, content: "hello?", role: "user" },
     ],
   },
-  output: [
-    { id: 4, content: "This is an AI Answer", role: ChatMessageRole.ai },
-  ],
+  output: [{ id: 4, content: "This is an AI Answer", role: "ai" }],
 };
 
 const defaultTemplate = {
   __type: "chat",
   messages: [
     {
       id: 0,
-      role: ChatMessageRole.system,
+      role: "system",
       content: "You are a chatbot",
     },
     {
       id: 1,
-      role: ChatMessageRole.user,
+      role: "user",
       content: "{{question}}",
     },
   ],

app/src/pages/playground/constants.tsx

@@ -1,8 +1,6 @@
-import { ChatMessageRole } from "@phoenix/store";
-
 export const NUM_MAX_PLAYGROUND_INSTANCES = 4;
 
-export const DEFAULT_CHAT_ROLE = ChatMessageRole.user;
+export const DEFAULT_CHAT_ROLE = "user";
 
 /**
  * Map of {@link ChatMessageRole} to potential role values.

app/src/pages/playground/playgroundUtils.ts

@@ -1,7 +1,6 @@
 import { PlaygroundInstance } from "@phoenix/store";
 import {
   ChatMessage,
-  ChatMessageRole,
   createPlaygroundInstance,
   generateMessageId,
 } from "@phoenix/store";

app/src/pages/playground/schemas.ts

@@ -6,7 +6,7 @@ import {
   SemanticAttributePrefixes,
 } from "@arizeai/openinference-semantic-conventions";
 
-import { ChatMessage, ChatMessageRole } from "@phoenix/store";
+import { ChatMessage } from "@phoenix/store";
 import { schemaForType } from "@phoenix/typeUtils";
 
 /**
@@ -76,7 +76,9 @@ export const outputSchema = z.object({
 /**
  * The zod schema for {@link chatMessageRoles}
  */
-export const chatMessageRolesSchema = z.nativeEnum(ChatMessageRole);
+export const chatMessageRolesSchema = schemaForType<ChatMessageRole>()(
+  z.enum(["user", "ai", "system", "tool"])
+);
 
 const chatMessageSchema = schemaForType<ChatMessage>()(
   z.object({

app/src/store/credentialsStore.tsx

@@ -0,0 +1,49 @@
+import { create, StateCreator } from "zustand";
+import { devtools, persist } from "zustand/middleware";
+
+export interface CredentialsProps {
+  /**
+   * The API key for the OpenAI API.
+   */
+  OPENAI_API_KEY?: string;
+  /**
+   * The API key for the Azure OpenAI API.
+   */
+  AZURE_OPENAI_API_KEY?: string;
+  /**
+   * The API key for the Anthropic API.
+   */
+  ANTHROPIC_API_KEY?: string;
+}
+
+export type CredentialKey = keyof CredentialsProps;
+
+export interface CredentialsState extends CredentialsProps {
+  /**
+   * Setter for a given credential
+   * @param credential the name of the credential to set
+   * @param value the value of the credential
+   */
+  setCredential: (params: {
+    credential: keyof CredentialsProps;
+    value: string;
+  }) => void;
+}
+
+export const createCredentialsStore = (
+  initialProps: Partial<CredentialsProps>
+) => {
+  const credentialsStore: StateCreator<CredentialsState> = (set) => ({
+    setCredential: ({ credential, value }) => {
+      set({ [credential]: value });
+    },
+    ...initialProps,
+  });
+  return create<CredentialsState>()(
+    persist(devtools(credentialsStore), {
+      name: "arize-phoenix-credentials",
+    })
+  );
+};
+
+export type CredentialsStore = ReturnType<typeof createCredentialsStore>;

app/src/store/index.tsx

@@ -1,3 +1,4 @@
 export * from "./pointCloudStore";
 export * from "./tracingStore";
-export * from "./playgroundStore";
+export * from "./playground";
+export * from "./credentialsStore";

app/src/store/playground/index.ts

@@ -0,0 +1,2 @@
+export * from "./playgroundStore";
+export * from "./types";