adjust UI to new schema

app/src/hooks/useChatMessageStyles.ts

@@ -6,6 +6,7 @@ export function useChatMessageStyles(
   role: string
 ): Pick<ViewStyleProps, "backgroundColor" | "borderColor"> {
   return useMemo<ViewStyleProps>(() => {
+    role = role.toLowerCase();
     if (role === "user" || role === "human") {
       return {
         backgroundColor: "grey-100",

app/src/pages/prompt/PromptChatMessages.tsx

@@ -1,18 +1,17 @@
-import React, { useMemo } from "react";
-import { useFragment } from "react-relay";
-import { graphql } from "relay-runtime";
+import React from "react";
+import { graphql, useFragment } from "react-relay";
 
 import { Flex, Text } from "@arizeai/components";
 
 import { TemplateLanguages } from "@phoenix/components/templateEditor/constants";
 import { TemplateLanguage } from "@phoenix/components/templateEditor/types";
 
 import {
+  PromptChatMessages__main$data,
   PromptChatMessages__main$key,
   PromptTemplateFormat,
 } from "./__generated__/PromptChatMessages__main.graphql";
 import { ChatTemplateMessage } from "./ChatTemplateMessage";
-import { PromptChatTemplate, PromptChatTemplateSchema } from "./schemas";
 
 const convertTemplateFormat = (
   templateFormat: PromptTemplateFormat
@@ -30,49 +29,69 @@ export function PromptChatMessages({
 }: {
   promptVersion: PromptChatMessages__main$key;
 }) {
-  const { template, templateType, templateFormat } = useFragment(
+  const { template, templateFormat } = useFragment(
     graphql`
       fragment PromptChatMessages__main on PromptVersion {
-        template
+        template {
+          __typename
+          ... on PromptChatTemplate {
+            messages {
+              ... on JSONPromptMessage {
+                role
+                jsonContent: content
+              }
+              ... on TextPromptMessage {
+                role
+                content
+              }
+            }
+          }
+          ... on PromptStringTemplate {
+            template
+          }
+        }
         templateType
         templateFormat
       }
     `,
     promptVersion
   );
 
-  if (templateType === "STRING") {
-    return <Text>{template}</Text>;
+  if (template.__typename === "PromptStringTemplate") {
+    return <Text>{template.template}</Text>;
+  }
+  if (template.__typename === "PromptChatTemplate") {
+    return (
+      <ChatMessages
+        template={template}
+        templateFormat={convertTemplateFormat(templateFormat)}
+      />
+    );
+  }
+  if (template.__typename === "%other") {
+    throw new Error("Unknown template type" + template.__typename);
   }
-
-  return (
-    <ChatMessages
-      template={template}
-      templateFormat={convertTemplateFormat(templateFormat)}
-    />
-  );
 }
 
 function ChatMessages({
   template,
   templateFormat,
 }: {
-  template: PromptChatTemplate | unknown;
+  template: Extract<
+    PromptChatMessages__main$data["template"],
+    { __typename: "PromptChatTemplate" }
+  >;
   templateFormat: TemplateLanguage;
 }) {
-  const messages = useMemo(() => {
-    const parsedTemplate = PromptChatTemplateSchema.safeParse(template);
-    if (!parsedTemplate.success) {
-      return [];
-    }
-    return parsedTemplate.data.messages;
-  }, [template]);
+  const { messages } = template;
   return (
     <Flex direction="column" gap="size-200">
       {messages.map((message, i) => (
+        // TODO: Handle JSON content for things like tool calls
         <ChatTemplateMessage
           key={i}
-          {...message}
+          role={message.role as string}
+          content={message.content || JSON.stringify(message.jsonContent)}
           templateFormat={templateFormat}
         />
       ))}