fix: normalize gemini roles (#7404)

Author: mikeldking

app/src/constants/generativeConstants.ts

@@ -45,3 +45,18 @@ export const AZURE_OPENAI_API_VERSIONS = [
   "2023-05-15",
   "2023-03-15-preview",
 ];
+
+/**
+ * Map of {@link ChatMessageRole} to potential role values.
+ * Used to map roles to a canonical role.
+ */
+export const ChatRoleMap: Record<ChatMessageRole, string[]> = {
+  user: ["user", "human"],
+  // Assistant used by OpenAI
+  // Model used by Gemini
+  ai: ["assistant", "bot", "ai", "model"],
+  // Developer is a newer tier of role from OpenAI
+  // While not 1 to 1 with system, it's a close match and so we treat it as such
+  system: ["system", "developer"],
+  tool: ["tool"],
+};

app/src/hooks/useChatMessageStyles.ts

@@ -1,23 +1,24 @@
 import { useMemo } from "react";
 
 import { ViewStyleProps } from "@phoenix/components/types";
+import { ChatRoleMap } from "@phoenix/constants/generativeConstants";
 
 export function useChatMessageStyles(
   role: string
 ): Pick<ViewStyleProps, "backgroundColor" | "borderColor"> {
   return useMemo<ViewStyleProps>(() => {
     const normalizedRole = role.toLowerCase();
-    if (normalizedRole === "user" || normalizedRole === "human") {
+    if (ChatRoleMap.user.includes(normalizedRole)) {
       return {
         backgroundColor: "grey-200",
         borderColor: "grey-500",
       };
-    } else if (normalizedRole === "assistant" || normalizedRole === "ai") {
+    } else if (ChatRoleMap.ai.includes(normalizedRole)) {
       return {
         backgroundColor: "blue-100",
         borderColor: "blue-700",
       };
-    } else if (normalizedRole === "system" || normalizedRole === "developer") {
+    } else if (ChatRoleMap.system.includes(normalizedRole)) {
       return {
         backgroundColor: "indigo-100",
         borderColor: "indigo-700",

app/src/pages/playground/constants.tsx

@@ -2,17 +2,6 @@ import { CanonicalParameterName } from "./__generated__/ModelSupportedParamsFetc
 
 export const NUM_MAX_PLAYGROUND_INSTANCES = 4;
 
-/**
- * Map of {@link ChatMessageRole} to potential role values.
- * Used to map roles to a canonical role.
- */
-export const ChatRoleMap: Record<ChatMessageRole, string[]> = {
-  user: ["user", "human"],
-  ai: ["assistant", "bot", "ai"],
-  system: ["system"],
-  tool: ["tool"],
-};
-
 /**
  * Parsing errors for parsing a span to a playground instance
  */

app/src/pages/playground/playgroundUtils.ts

@@ -6,6 +6,7 @@ import { TemplateFormats } from "@phoenix/components/templateEditor/constants";
 import { getTemplateFormatUtils } from "@phoenix/components/templateEditor/templateEditorUtils";
 import { TemplateFormat } from "@phoenix/components/templateEditor/types";
 import {
+  ChatRoleMap,
   DEFAULT_CHAT_ROLE,
   DEFAULT_MODEL_PROVIDER,
 } from "@phoenix/constants/generativeConstants";
@@ -51,7 +52,6 @@ import {
   InvocationParameterInput,
 } from "./__generated__/PlaygroundOutputSubscription.graphql";
 import {
-  ChatRoleMap,
   INPUT_MESSAGES_PARSING_ERROR,
   MODEL_CONFIG_PARSING_ERROR,
   MODEL_CONFIG_WITH_INVOCATION_PARAMETERS_PARSING_ERROR,