feat(playground): Implement message role picker callback (#4909)

* feat(playground): Implement message role picker callback

* style(playground): Shorten import path for playgroundUtils

app/src/pages/playground/MessageRolePicker.tsx

@@ -5,6 +5,8 @@ import { Item, Picker } from "@arizeai/components";
 
 import { ChatMessageRole } from "@phoenix/store";
 
+import { isChatMessageRole } from "./playgroundUtils";
+
 const hiddenLabelCSS = css`
   .ac-field-label {
     display: none;
@@ -22,22 +24,30 @@ type MessageRolePickerProps = {
    * @default true
    */
   includeLabel?: boolean;
+  /**
+   * Callback for when the message role changes
+   */
+  onChange: (role: ChatMessageRole) => void;
 };
 
 export function MessageRolePicker({
   role,
   includeLabel = true,
+  onChange,
 }: MessageRolePickerProps) {
   return (
     <Picker
       selectedKey={role}
       css={!includeLabel ? hiddenLabelCSS : undefined}
       label="Role"
-      data-testid="inferences-time-range"
-      aria-label={`Time range for the primary inferences`}
+      data-testid="messages-role-picker"
+      aria-label={`Role for the chat message`}
       size="compact"
-      onSelectionChange={() => {
-        // TODO: fill out
+      onSelectionChange={(e) => {
+        if (!isChatMessageRole(e)) {
+          throw new Error(`Invalid chat message role: ${e}`);
+        }
+        onChange(e);
       }}
     >
       <Item key="system">System</Item>

app/src/pages/playground/PlaygroundChatTemplate.tsx

@@ -37,7 +37,23 @@ export function PlaygroundChatTemplate(props: PlaygroundChatTemplateProps) {
           <li key={index}>
             <Card
               title={
-                <MessageRolePicker includeLabel={false} role={message.role} />
+                <MessageRolePicker
+                  includeLabel={false}
+                  role={message.role}
+                  onChange={(role) => {
+                    updateInstance({
+                      instanceId: id,
+                      patch: {
+                        template: {
+                          __type: "chat",
+                          messages: template.messages.map((message, i) =>
+                            i === index ? { ...message, role } : message
+                          ),
+                        },
+                      },
+                    });
+                  }}
+                />
               }
               variant="compact"
               backgroundColor="light"

app/src/pages/playground/playgroundUtils.ts

@@ -0,0 +1,11 @@
+import {
+  ChatMessageRole,
+  chatMessageRoles,
+} from "@phoenix/store/playgroundStore";
+
+/**
+ * Checks if a string is a valid chat message role
+ */
+export function isChatMessageRole(role: unknown): role is ChatMessageRole {
+  return chatMessageRoles.includes(role as ChatMessageRole);
+}

app/src/store/playgroundStore.tsx

@@ -19,10 +19,15 @@ export type PlaygroundTemplate =
   | PlaygroundChatTemplate
   | PlaygroundTextCompletionTemplate;
 
+/**
+ * Array of roles for a chat message with a LLM
+ */
+export const chatMessageRoles = ["user", "ai", "system", "tool"] as const;
+
 /**
  * The role of a chat message with a LLM
  */
-export type ChatMessageRole = "user" | "ai" | "system" | "tool";
+export type ChatMessageRole = (typeof chatMessageRoles)[number];
 
 /**
  * A chat message with a role and content