From b2a8bd0245bea2cc39e335dabaa70a856cacd423 Mon Sep 17 00:00:00 2001
From: Anthony Powell <apowell@arize.com>
Date: Tue, 8 Oct 2024 11:14:57 -0400
Subject: [PATCH] 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 | 18 ++++++++++++++----
 .../playground/PlaygroundChatTemplate.tsx      | 18 +++++++++++++++++-
 app/src/pages/playground/playgroundUtils.ts    | 11 +++++++++++
 app/src/store/playgroundStore.tsx              |  7 ++++++-
 4 files changed, 48 insertions(+), 6 deletions(-)
 create mode 100644 app/src/pages/playground/playgroundUtils.ts

diff --git a/app/src/pages/playground/MessageRolePicker.tsx b/app/src/pages/playground/MessageRolePicker.tsx
index ee44b457e1..d10e46fbc2 100644
--- a/app/src/pages/playground/MessageRolePicker.tsx
+++ b/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>
diff --git a/app/src/pages/playground/PlaygroundChatTemplate.tsx b/app/src/pages/playground/PlaygroundChatTemplate.tsx
index f2103de987..7263eb7b47 100644
--- a/app/src/pages/playground/PlaygroundChatTemplate.tsx
+++ b/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"
diff --git a/app/src/pages/playground/playgroundUtils.ts b/app/src/pages/playground/playgroundUtils.ts
new file mode 100644
index 0000000000..50a154f384
--- /dev/null
+++ b/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);
+}
diff --git a/app/src/store/playgroundStore.tsx b/app/src/store/playgroundStore.tsx
index 3b47af994e..8315720155 100644
--- a/app/src/store/playgroundStore.tsx
+++ b/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